[LeetCode] 1814. Count Nice Pairs in an Array

문제 출처: https://leetcode.com/problems/count-nice-pairs-in-an-array/

문제

음이 아닌 정수로 구성된 배열 nums가 주어진다. rev(x)를 음이 아닌 정수 x를 뒤집는 것으로 정의하자. 예를 들어, rev(123) = 321, 이고 rev(120) = 21이다. 지시자의 쌍 (i, j)가 다음을 따를 때 나이스(nice)하다:

 

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

 

나이스한 지시자의 쌍의 수 를 반환하라. 숫자가 매우 클 경우를 대비하여, 10^9 + 7로 나눈 나머지를 반환하라.

예제

Input: nums = [42,11,1,97]
Output: 2

Explanation: The two pairs are:
 - (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.

Input: nums = [13,10,35,24,76]
Output: 4

풀이

나이스한 지시자의 쌍 (i, j)의 조건 중, 핵심이 되는 nums[i] + rev(nums[j]) == nums[j] + rev(nums[i])을 다음과 같이 표현할 수 있다.

• nums[i] - rev(nums[i]) == nums[j] - rev(nums[j])

 

좌항과 우항에 각각 하나의 변수(지시자)만 있도록 만든 것이다.

 

기존의 수식을 변형하지 않고 풀이를 한다면, 어떤 i와 i보다 큰 모든 j에 대한 nums[i] + rev(nums[j]) 그리고 nums[j] + rev(nums[i]) 계산이 필요하다. 어림잡아 n^2번의 계산이 수행되어야 한다.

 

하지만 위와 같이 변경을 하게 되면, 모든 인덱스에 대해서 한 번의 계산만 필요하다.

const arr = [];
nums.forEach(val => arr.push(val - rev(val)));

이후, 나이스한 관계의 쌍을 찾기 위해 다음과 같이 반복문을 구성할 수 있다.

let ans = 0;
for (let i = 0; i < arr.length-1; i++) {
    for (let j = i+1; j < arr.length; j++) {
        if (arr[i] == arr[j]) { ans += 1; }
    }
}

이 구조에서의 시간 복잡도를 생각해보면 O(n^2)이 될 것이다. 기껏 위에서 n^2의 반복을 하지 않으려 수식을 변경했는데, 여기서 n^2의 시간이 소요되는 불상사가 생긴 것이다.

 

이 상황을 타개하기 위해서 arr을 동일 값을 갖는 인덱스 별로 나누어보자.

arr = [1, 2, 3, 2, 2, 3, 1]

arr[i] == 1 인 인덱스: [0, 6]
arr[i] == 2 인 인덱스: [1, 3, 4]
arr[i] == 3 인 인덱스: [2, 5]

위의 예제 상황에서 각각 몇 개의 나이스한 쌍을 만들 수 있는가? 1개, 3개, 1개의 쌍을 만들 수 있고 이들의 합이 최종 얻고자 하는 값이 된다.

 

몇 번의 수기 작업을 통해서 어떤 규칙이 존재하는지에 대해 알아보면,

1) 3개의 인덱스 a, b, c가 동일한 값을 같는 원소라고 가정 ( a < b < c )

    i < j 관계: a, b / a, c / b, c : 3

1) 4개의 인덱스 a, b, c, e가 동일한 값을 같는 원소라고 가정 ( a < b < c < e )

    i < j 관계: a, b / a, c / b, c + a, d / b, d / c, d : 3 + 3 = 6

2) 5개의 인덱스 a, b, c, e, f가 동일한 값을 같는 원소라고 가정 ( a < b < c < e < f )

    i < j 관계: a, b / a, c / b, c / a, d / b, d / c, d + a, f / b, f / c, f / e, f : 6 + 4 = 10

인덱스의 수에 따라 값이 규칙적으로 변하는 것이 눈에 띈다. :

• memo[x] = memo[x-1] + x - 1 (memo: i < j 쌍의 수를 저장하는 배열, x: 동일 값을 갖는 인덱스의 개수)

 

이 규칙을 이용하여 코드로 작성하면 다음과 같다.

var countNicePairs = function(nums) {
    // hash => key:nums[i] - rev(nums[i]), value: 개수
    let hash = new Map();
    nums.forEach(val => {
        const n = val - rev(val);
        if (!hash.has(n)) { hash.set(n, 0); }
        hash.set(n, hash.get(n) + 1);
    });

    // pairsNum: 'x'개의 동일 값이 있을 때, 만들 수 있는 i < j 쌍의 수를 저장하는 배열 
    let pairsNum = [0, 0, 1, 3, 6, 10];

    function getPairsNum(x) {
        if (pairsNum[x]) { return pairsNum[x]; }
        if (x == 0 || x == 1) return 0;

        pairsNum[x] = getPairsNum(x-1) + x - 1;
        return pairsNum[x];
    }

    // hash를 순회하면서 최종 답 계산
    let ans = 0;
    hash.forEach(val => {
        ans += getPairsNum(val);
        ans %= 1000000007;
    })

    return ans;
};

Github: 1814.js

opwe37/Algorithm-Study