LeetCode题目笔记——1814. 统计一个数组中好对子的数目

题目描述

给你一个数组 nums ，数组中只包含非负整数。定义 rev(x) 的值为将整数 x 各个数字位反转得到的结果。比方说 rev(123) = 321 ， rev(120) = 21 。我们称满足下面条件的下标对 (i, j) 是 好的 ：

• 0 <= i < j < nums.length
• nums[i] + rev(nums[j]) == nums[j] + rev(nums[i])
  请你返回好下标对的数目。由于结果可能会很大，请将结果对 10⁹ + 7 取余 后返回。

示例1：

输入：nums = [42,11,1,97]
输出：2
解释：两个坐标对为：
(0,3)：42 + rev(97) = 42 + 79 = 121, 97 + rev(42) = 97 + 24 = 121 。
(1,2)：11 + rev(1) = 11 + 1 = 12, 1 + rev(11) = 1 + 11 = 12 。

示例2：

输入：nums = [13,10,35,24,76]
输出：4

提示：

• 1 <= nums.length <= 10⁵
• 0 <= nums[i] <= 10⁹

题目难度——中等

方法一：模拟（超时）（参考）

  根据根据题意，我们可以按照他的要求模拟这一过程，首先新建一个数组，然后求每个数的反转数，再用两个循环找满足要求的下标对。根据提示最多有100000个数据，因此肯定会超时。在本地测试10⁵个数据的话，需要10多秒钟。

代码/Python

class Solution:
    def countNicePairs(self, nums: List[int]) -> int:
        n = len(nums)
        revNums = [0] * n
        for i in range(n):
            tmp = nums[i]
            x = 0
            while tmp != 0:
                x = x * 10 + tmp % 10
                tmp //= 10
            revNums[i] = x
        
        res = 0
        for i in range(n):
            for j in range(i + 1, n):
                if nums[i] - revNums[i] == nums[j] - revNums[j]:
                    res += 1
        
        return res

方法二：优化

  再看一下题目中的等式，我们合并一下同类项，把i和j都移到同一侧，就变成了nums[i] - rev(nums[i]) == nums[j] - rev(nums[j])，说明我们可以直接在原数组上对每个元素减去它本身的反数。再看超时问题，我们需要很快的访问，所以可以用一个哈希表来记录每个数字出现的次数，根据规律，出现两次则下标对就是1，三次就是2，以此类推，所以我们让答案加上counter[x] - 1.

代码

class Solution:
    def countNicePairs(self, nums: List[int]) -> int:
        size = len(nums)
        for i, x in enumerate(nums):
            tmp = 0
            while x != 0:
                tmp = tmp * 10 + x % 10
                x //= 10
            nums[i] -= tmp
        
        res = 0
        counter = defaultdict(int)
        for i, x in enumerate(nums):
            counter[x] += 1
            res += counter[x] - 1

        return res % (10 ** 9 + 7)

  还是很慢，还可以优化，我们没必要每次counter[x]加一后让答案加一次，因为这样对某个counter[x]来说，他会是加一个等差数列，我们直接用等差数列的求和公式。

class Solution:
    def countNicePairs(self, nums: List[int]) -> int:
        size = len(nums)
        for i, x in enumerate(nums):
            tmp = 0
            while x != 0:
                tmp = tmp * 10 + x % 10
                x //= 10
            nums[i] -= tmp
        
        res = 0
        counter = defaultdict(int)
        for i, x in enumerate(nums):
            counter[x] += 1
        
        for v in counter.values():
            res += (1 + v - 1) * (v - 1) // 2

        return res % (10 ** 9 + 7)

  还是慢，我们还可以继续优化。有三次循环，我们可以想办法在一个循环里解决。回到原题目，左右相等，其实我们可以根本不管他们具体的值是多少，只用关心它们的出现次数就好，还是用哈希表，对每个x，记录x-rev(x)的出现次数，然后像上面那样更新答案。

class Solution:
    def countNicePairs(self, nums: List[int]) -> int:
        size = len(nums)
        occurs = defaultdict(int)
        res = 0
        for i, x in enumerate(nums):
            rev = int(str(x)[::-1])
            res += occurs[x - rev]
	        occurs[x - rev] += 1

        return res % (10 ** 9 + 7)

总结

  暴力的复杂度是O(N²)，优化后能达到O(N)。用到了字典，所以空间是O(N)。