[LeetCode] 923. 3Sum With Multiplicity

https://leetcode.com/problems/3sum-with-multiplicity/

3Sum With Multiplicity - LeetCode

 

1. Problem 

Given an integer array A, and an integer target, return the number of tuples i, j, k  such that i < j < k and A[i] + A[j] + A[k] == target.

As the answer can be very large, return it modulo 10^9 + 7.

Note:

1. 3 <= A.length <= 3000
2. 0 <= A[i] <= 100
3. 0 <= target <= 300

2. Algorithm

결국 이 문제는 길이가 3이면서, 원소들의 합이 target 값이 되는 부분배열의 갯수를 찾아주면 된다. (단, A[i] == A[j] 이더라도 i != j이면 중복이라고 보지 않음)

문제에서 중요한 제약 조건은, 0 <= A[i] <= 100 으로 배열 내의 값이  0 ~ 100까지 총 101개만 존재한다는 사실이다. 이 사실을 이용하면 O(N + 101*101)만에 풀 수가 있다. 따라서 배열 내의 값 범위가 K로 유한하게 정해져있다면, 범위에 따라서 우리는 O(N+K^2)에 문제를 풀 수 있게 된다.

 

Step 1. 각 원소들의 출현 횟수를 카운트배열에 저장한다.

Step 2. 이제 각 원소들을 가지고 조합을 하며 가능한 갯수들을 res에 저장해줄 것이다. i + j + k = target이 되는 원소들의 조합 케이스는 아래와 같다. (해당 원소들의 정렬 순서는 신경쓰지 않아도 된다. 덧셈은 순서에 상관없기 때문이다.)

첫번째, i != j != k 인 경우, 그대로 cnt[i] * cnt[j] * cnt[k]를 하면 만들어 낼 수 있는 조합의 갯수가 된다.

두번째, i == j != k 인 경우, i == j 인 수를 2번 뽑아야지 조합이 가능하므로 cnt[i]의 갯수가 2 이상일 때만 조합을 진행한다. 조합은 cnt[i]개 중에서 2개를 뽑아야 하므로, cnt[i] * (cnt[j] -1) / 2 * cnt[k] 가 된다. 

세번째,  i != j == k 인 경우, j == k 인 수를 2번 뽑아야지 조합이 가능하므로 cnt[j]의 갯수가 2 이상일 때만 조합을 진행한다. 조합은 cnt[j]개 중에서 2개를 뽑아야 하므로, cnt[i] * cnt[j] * (cnt[k] -1) / (2*1)  가 된다. 

네번째, i == j == k 인 경우, i == j == k 인 수를 3번 뽑아야지 ㅈ합이 가능하므로 cnt[k]의 갯수가 3 이상일 때만 조합을 진행한다. 조합은 cnt[k]개 중에서 3개를 뽑아야 하므로, cnt[i] * (cnt[j]-1) * (cnt[k] - 2) / (3*2*1) 이 된다.

 

아래의 코드에서는 조합의 갯수가 Integer 범위를 넘는 경우가 있기 때문에, 문제에서 제시한 것과 같이 결과 값마다 나머지 연산을 진행하며 계산을 해주었다. 범위를 넘는 케이스를 정확하게 계산하기 위해 double 타입의 temp 변수를 두어 진행하는 부분의 코드가 깔끔하지는 않지만, 기본적인 로직은 위의 케이스와 같다.

 

class Solution {
    public int threeSumMulti(int[] A, int target) {
        int cnt[] = new int[101];
        int res = 0, mod = (int)Math.pow(10,9)+7;
        //Step 1
        for(int n : A) cnt[n]++;
        for(int i = target > 100 ? 100 : target; i >= 0; i--) {
            if(cnt[i] == 0) continue;
            for(int j = target - i > 100 ? 100 : target - i; j >= 0; j--) {
                int k = target-i-j > 100 ? 100 : target - i - j;
                if(cnt[j] == 0 || cnt[k] == 0 || !(i >= j && j >= k)) continue;
                if(i != j && j != k) res+=(cnt[i]*cnt[j]*cnt[k]);
                else if(i == j && j != k && cnt[i] >= 2) res+=(cnt[i]*(cnt[j]-1)/2*cnt[k])%mod;
                else if(i != j && j == k && cnt[j] >= 2) res+=(cnt[i]*cnt[j]*(cnt[k]-1)/2)%mod;
                else if(i == j && j == k && cnt[k] >= 3) {
                    double temp = (cnt[i]/6.0 %mod)*((cnt[k]-1)%mod)*((cnt[j]-2)%mod);
                    res += (temp)%mod;
                }
            }
        }
        return res%mod;
    }
}

 

3. retrospective

수학 개념을 응용하는 류의 문제는 자신이 없었는데, 그래도 잘 푼 편에 속했던 것 같다. 문제를 풀면서 블로그에 일일이 업데이트하는 것이 힘들어서 스킵해왔는데, 이 문제는 어떻게 풀었는지 기억하고 싶어서 남긴다.