统计一个数组中好对子的数目

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])
请你返回好下标对的数目。由于结果可能会很大，请将结果对 109 + 7 取余 后返回。

提示：

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/count-nice-pairs-in-an-array
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

解析

分析

先得到nums对应的rev数组。根据等式nums[i] + rev(nums[j]) == nums[j] + rev(nums[i])可以得到nums[i] - rev(nums[i]) == nums[j] - rev(nums[j])。所以我们求一个diff数组，其中diff[i]=nums[i]-rev[i]。

问题转化为：
找到i，j对，i<j，且diff[i]==diff[j]。假设diff为val的有k个，则这样的对数有k*(k-1)/2个。对所有不同diff的对数求和，就能得到所有的满足条件的下标对数。

这道题需要注意的是溢出问题。由于n<=10^{5，所以k*(k-1)可能是long。但是rev和diff是不会溢出，因为nums[i]<109,而int最大超过2x1010。}

代码

class Solution {
    int MOD=1000000007;
    public int countNicePairs(int[] nums) {
        int[] revs=getRevs(nums);
        // System.out.println(Arrays.toString(revs)); 
        int[] diff=getDiff(nums,revs);
        // System.out.println(Arrays.toString(diff)); 
        Map<Integer,Integer> frq=groupByDiff(diff);
        
        // System.out.println(frq); 
        return getCount(frq);
    }
    
    private int getCount(Map<Integer,Integer> frq){
        int sum=0;
        for(int f:frq.values()){
            long k=f;
            sum+=((k*(k-1))>>1)%MOD;
            sum%=MOD;
        }      
        return sum;
    }
    
    private Map<Integer,Integer> groupByDiff(int[] diff){
        Map<Integer,Integer> res=new HashMap<Integer,Integer>();
        for(int val:diff){
            int cnt=0;
            if(res.containsKey(val)){
                cnt=res.get(val);
            }
            res.put(val,cnt+1);
        }
        return res;
    }
    private int[] getDiff(int[] nums,int[] revs){
        int n=nums.length;
        int[] res=new int[n];
        for(int i=0;i<n;i++){
            res[i]=nums[i]-revs[i];//溢出？
        }
        return res;
    }
    
    private int[] getRevs(int[] nums){
        int n=nums.length;
        int[] res=new int[n];
        for(int i=0;i<n;i++){
            res[i]=getR(nums[i]);
        }
        return res;
    }
    
    private int getR(int val){
        int res=0;
        while(val!=0){
            res*=10;
            int mod=val%10;
            val/=10;
            res+=mod;
        }
        return res;
    }
}

复杂度分析

空间：O(N)用于放数组和map（其实可以优化成diff就用原来的rev空间）
时间：O(N)。

优秀题解

得到diff数组，按照从小到大排序。双指针i和j，j一直走到和i位置的值不相等，求得k并求和，然后更新i=j。

核心代码：

Arrays.sort(revNums);
int ans = 0;
int mod = 1_000_000_007;
int i = 0;
 for (int j = 0; j < revNums.length; ++j) {
    if (revNums[j] != revNums[i]) {
        long n = (long)(j - i) * (long)(j - i - 1) >> 1;
        n %= mod;
        ans += (int)n;
        if (ans > mod) ans -= mod;
        i = j;
    }
}
if (i != revNums.length - 1) {
    int j = revNums.length;
    long n = (long)(j - i) * (long)(j - i - 1) >> 1;
    n %= mod;
    ans += (int)n;
    if (ans > mod) ans -= mod;
}

时间：O(NlogN)
空间：O(N)放rev数组