Program to count nice pairs in an array in Python

Suppose we have an array called nums with non-negative values. We have to find the number of nice pairs of indices present in the array. If the answer is too large, then return answer mod 10^9+7. Here a pair of indices (i, j) is said to be nice when they satisfy all of these conditions:

1. 0 ≤ i < j < size of nums
2. nums[i] + rev(nums[j]) is same as nums[j] + rev(nums[i])

Note: Here rev() reverses only the positive part of an integer, so if rev(564) is there, it means 465, but if rev(540) is there, it returns 45.

Example

If the input is nums = [97, 2, 42, 11], then the output will be 2 because there are two pairs (0,2) and (1,3):

• For pair (0,2): 97 + rev(42) = 97 + 24 = 121, and 42 + rev(97) = 42 + 79 = 121
• For pair (1,3): 2 + rev(11) = 2 + 11 = 13, and 11 + rev(2) = 11 + 2 = 13

Algorithm

To solve this, we will follow these steps −

1. Create a dictionary to store the frequency of (num - rev(num)) values
2. For each number, calculate its reverse and store num - rev(num) in the dictionary
3. For each frequency count, calculate the number of pairs using the formula n * (n-1) / 2
4. Return the total count modulo 10^9+7

Implementation

from collections import defaultdict

def solve(nums):
    m = (10**9) + 7
    dic = defaultdict(int)
    
    # Calculate num - rev(num) for each number
    for num in nums:
        rev = int(str(num)[::-1])
        dic[num - rev] += 1
    
    res = 0
    # Count pairs for each group
    for val in dic.values():
        res += (val * (val - 1)) // 2
    
    return res % m

# Test the function
nums = [97, 2, 42, 11]
print(solve(nums))

2

How It Works

The key insight is that for indices i and j to form a nice pair, we need:

nums[i] + rev(nums[j]) = nums[j] + rev(nums[i])

Rearranging: nums[i] - rev(nums[i]) = nums[j] - rev(nums[j])

So we group numbers by their (num - rev(num)) value. Within each group, any two numbers can form a nice pair. For n numbers in a group, the number of pairs is n * (n-1) / 2.

Conclusion

This solution efficiently finds nice pairs by grouping numbers with the same (num - rev(num)) difference. The time complexity is O(n) and space complexity is O(n), making it optimal for large arrays.