[LeetCode] 1814. Count Nice Pairs in an Array

https://leetcode.com/problems/count-nice-pairs-in-an-array/description/

Count Nice Pairs in an Array - LeetCode

leetcode - Count Nice Pairs in an Array

문제 유형: 배열, 해시, 수학

문제 난이도: Medium

 

문제

You are given an array nums that consists of non-negative integers. Let us define rev(x) as the reverse of the non-negative integer x. For example, rev(123) = 321, and rev(120) = 21. A pair of indices (i, j) is nice if it satisfies all of the following conditions:

• 0 <= i < j < nums.length
• nums[i] + rev(nums[j]) == nums[j] + rev(nums[i])

Return the number of nice pairs of indices. Since that number can be too large, return it modulo 10^9 + 7.

 

당신은 자연수 배열 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])

좋은 쌍의 개수를 반환하시오. 값이 클 수 있으니, 1e9+7로 나눈 나머지를 반환하시오.

 

 

풀이

단순히 투 포인터로 두 쌍을 계산해보면서 구하는 방식은 O(n^2)로, 시간초과가 난다.

 

시간초과가 나지 않는 방법은 다음과 같다:

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

nums[i] - rev(nums[i]) = nums[j] -rev([nums[j])로 바꿀 수 있다.

선형적으로 배열을 탐색하면서 해시 테이블의 key {nums[i]-rev(nums[i])}에다가 value로 1을 누적시킨다.

 

모든 순회가 끝나면, 각 key별로 만들 수 있는 쌍의 개수 = (value * (value-1) / 2)를 누적시킨다.

 

 

코드

C++

class Solution {
public:
    long long rev(int num) {
        string str = "";
        while(num > 0) {
            str += to_string(num%10);
            num /= 10;
        }
        if(str.size() == 0) str = "0";
        return stoll(str);
    }

    int countNicePairs(vector<int>& nums) {
        const int n = nums.size();
        ios_base::sync_with_stdio(0);
        cin.tie(0);
        cout.tie(0);
        
        unordered_map<long, long> um;
        long cnt = 0;

        for(int i = 0; i < n; i++) {
            long diff = nums[i] - rev(nums[i]);
            um[diff]++;
        }

        for(const auto &map : um) {
            cnt += map.second * (map.second - 1)/2;
        }

        return cnt%(int)(1e9+7);
    }
};