[Leetcode] 2179. Count Good Triplets in an Array

문제는 여기서 확인할 수 있다.

꽤 어려운 문제였고 고백하자면 Hint 3개를 다 보고 나서야 풀 수 있었다.

해외 취업의 길은 멀고도 험하다.

거기 끝에서 뭐하시는지?

문제 설명 - Count Good Triplets in Two Arrays

두 개의 0-indexed 배열 nums1, nums2가 주어진다.
이 두 배열은 모두 [0, 1, ..., n-1]의 순열이며, 길이는 n이다.

여기서 좋은 삼중 조합(good triplet) 은 다음 조건을 만족하는 (x, y, z) 세 값의 조합이다:

• x, y, z는 모두 다른 값이며
• x, y, z의 값 순서대로,
  • nums1에서의 위치: pos1[x] < pos1[y] < pos1[z]
  • nums2에서의 위치: pos2[x] < pos2[y] < pos2[z]

즉, 두 배열 모두에서 x, y, z가 증가하는 순서로 등장해야 한다.

목표는 가능한 모든 good triplet의 개수를 세는 것이다.

 

예시

예제 1

Input: nums1 = [2, 0, 1, 3]
       nums2 = [0, 1, 2, 3]
Output: 1

가능한 조합은 총 4개다:

• (2,0,1), (2,0,3), (2,1,3), (0,1,3)

이 중에서 두 배열에서 모두 순서가 유지되는 조합은 (0,1,3) 하나뿐이다.

따라서 정답은 1이다.

예제 2

Input: nums1 = [4,0,1,3,2]
       nums2 = [4,1,0,2,3]
Output: 4

good triplet으로 인정되는 조합은 (4,0,3), (4,0,2), (4,1,3), (4,1,2) 4가지다.

 

조건 정리

• n == nums1.length == nums2.length
• 3 <= n <= 10⁵
• 각 배열은 [0, 1, ..., n-1]의 순열
• 즉, 모든 원소는 유일하며 서로 다름

 

풀이 전략

이 문제는 단순한 브루트포스로는 풀 수 없다.
모든 삼중 조합을 순회하면 O(n³)이 되기 때문이다.
대신 좌측에 있는 작은 값의 개수와 우측에 있는 큰 값의 개수를 잘 세는 방식으로 최적화할 수 있다.

 

그런데, 좌측에 있는 작은 값의 개수를 알면 우측에 있는 큰 값의 개수를 알 수 있다.

$S_{2}$ 집합은 $i$ 번째 위치보다 왼쪽에 있으면서 $V_{i}$ 보다 작은 값을 가진다.

이것만 있으면 $S_{3}$ 집합 역시 쉽게 구할 수 있다.

그리고 나면 결국 구하고자 하는 값인 $S_{2} x S_{3}$ 을 쉽게 계산할 수 있다.

S1 집합은 i 개 원소 중 S2 집합을 제외한 것이다.

 

그렇기 때문에 결국, 특정 위치에서 내 왼쪽으로 나보다 작은 개수를 빠르게 셀 수 있으면 된다.

이 값은 Segment Tree 또는 Lower bound 를 통해 구할 수 있다.

나는 Lower Bound를 통해 구했다.

class Solution:
    def goodTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        total = len(nums1)
        idx_dict = {}
        for idx, num in enumerate(nums2):
            idx_dict[num] = idx
        
        arr = []
        for num in nums1:
            arr.append(idx_dict[num])
        
        small_cnt = {}
        sorted_list = []
        for idx, num in enumerate(arr):
            l, r = 0, len(sorted_list)

            while l < r:
                m = (l + r) // 2

                if sorted_list[m] < num:
                    l = m + 1
                else:
                    r = m
            
            small_cnt[idx] = l
            sorted_list.insert(l, num)
        
        ans = 0
        for idx, num in enumerate(arr):
            ans += small_cnt[idx] * (total - num - (idx + 1 - small_cnt[idx]))
        
        return ans

 

핵심 아이디어

핵심 아이디어는 아래와 같다. 

1. 먼저 nums2의 각 값의 위치를 빠르게 찾을 수 있도록 pos2[val] 맵을 만든다.
2. nums1을 순회하면서, 현재 값 v가 nums2에서 어느 위치에 있는지 확인한다.
3. 지금까지 본 값 중에서 nums2에서 v보다 왼쪽에 있는 값들의 개수를 세고,
4. 앞으로 나올 값 중에서 nums2에서 v보다 오른쪽에 있는 값들의 개수를 세면,
5. 이 조합을 기준으로 good triplet의 수를 계산할 수 있다.

이 과정을 효율적으로 하기 위해서 Binary Indexed Tree (Fenwick Tree) 또는 Segment Tree를 사용할 수 있다.

난 구현 복잡도를 줄이기 위해 Lower Bound를 사용했지만 이게 더 빠를 것 같은 느낌이 든다.

• 앞에서부터 순회하며 left_count를 기록하고
• 뒤에서부터 순회하며 right_count를 누적해 계산한다
• 각 v에 대해 left_count[v] * right_count[v]를 더하면 총 triplet 수가 된다.

 

이 문제는 순열과 좌우 비교가 결합된 문제로, 세그먼트 트리나 펜윅 트리를 연습하기에 매우 적합하다.
트리 구조 없이도 정렬된 배열을 활용한 좌우 카운팅 기법으로도 풀 수 있으므로, 다양한 풀이법을 연습해보는 걸 추천한다.

아래는 힌트이다.