Find K Pairs With Smallest Sums - 리트코드

https://leetcode.com/problems/find-k-pairs-with-smallest-sums

어떻게든 공간복잡도를 줄여보겠다는 혼신의 노력....

평가

결과 : 실패
풀이 시간 : 2시간 이상 고민

다른 분들의 풀이들을 봐도 이해가 잘 안가 코드를 그대로 따라치면서 이해를 해보았습니다.
저 역시 풀지 못한 문제였고, 다른 분들의 풀이를 통해 이해할 수 있었는데요.
설명이 제대로 나온 블로그를 찾지 못해 이해하는데 약간 어려움을 겪었습니다.
따라서 제가 푼 건 아니지만, 제가 이해한 내용을 바탕으로 풀이를 설명하겠습니다.

아이디어

해당 문제를 해결하기 위해 Heap을 사용해야 합니다. 이 부분에 대한 설명은 생략하겠습니다.
다만 이 문제를 풀어보면 공간복잡도에서 걸리는데요. 모든 Case를 고려할 수 없다는 의미가 됩니다.

문제에서는 두 원소의 합이 작은 K개의 Case를 골라야 하는데요.
모든 Case를 보지 않고 K개의 Case를 어떻게 찾을 수 있을까요?

다음과 같은 방식을 사용합니다.

[1,7,11] 과 [2,4,6]이라는 두 배열이 있습니다.

먼저 두 번째 배열의 값을 고정한 상태로 조합을 만들어 Heap에 넣습니다.

Heap에서 가장 상단에 있는 [1,2] 를 poll합니다.
그러고 빠져 나온 값을 기준으로 두 번째 배열을 한 칸 이동시킨 값을 Heap에 추가합니다. [1,2] -> [1,4]

이제 Heap은 다음과 같은 값을 가지게 됩니다.

Heap에서 가장 상단에 있는 [1,4]를 뺀 뒤 마찬가지로 두 번째 배열을 한 칸 이동시켜 Heap에 추가합니다. [1,4] -> [1,6]

다음은 [1,6]이 빠지게 되는데요. 두 번째 배열에서 6 이후에는 값이 존재하지 않습니다. 따라서 더 이상 다음 경우를 추가하지 않습니다.

해당 아이디어의 핵심은 두 배열이 정렬되어 있다는 것입니다.
처음 선택되는 [1,2] 는 합이 가장 작은 조합입니다.
그 다음으로 합이 작으려면 [1,4] 거나 [7,2] 둘 중 하나에 위치해야 합니다.
[1,4]와 [7,2] 중 어떤 게 더 합이 작은지는 Heap에게 맡기고 우리의 역할은 두 값을 Heap에 넣어주는 것입니다.


이러한 방식으로 전체 Case를 넣지 않고도 K번째로 작은 조합을 판단할 수 있게 됩니다.

정답 코드

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // nums1과 nums2를 조합해서 합이 작은 순서대로 정렬. k개를 반환한다
        Queue<int[]> queue = new PriorityQueue<>(new Comparator<>() {
            public int compare(int[] n1, int[] n2) {
                return n1[0] + n1[1] - n2[0] - n2[1];
            }
        });

        // 두 번째 배열의 첫 번째 원소만을 사용해 조합을 만든다
        // ex. [0번째, 0번째], [1번째, 0번째]. [2번째, 0번째]
        for(int i=0; i<Math.min(nums1.length, k);i++) {
            queue.offer(new int[] {nums1[i], nums2[0], 0});
        }

        // while(!queue.isEmpty()) {
        //     System.out.println(queue.poll());
        // }
        List<List<Integer>> answers = new ArrayList<>();
        int idx =-1;
        while(!queue.isEmpty()) {
            idx++;
            if (idx >= k) {
                break;
            }
            
            int[] node = queue.poll();
            answers.add(List.of(node[0], node[1]));
            
            if (node[2] == nums2.length - 1) {
                continue;
            }

            // [0번째,0번째] 원소를 확인했다면 -> [0번째, 1번째] 원소를 추가한다
            // [0번째,1번째] 원소를 확인했다면 -> [0번째, 2번째] 원소를 추가한다
            queue.offer(new int[] {node[0], nums2[node[2]+1], node[2]+1});
        }
        return answers;
    }

}