[Java] 373. Find K Pairs with Smallest Sums

문제 파악

비 내림차순인 두 정수 배열 nums1 과 nums2 그리고 k가 주어진다.

k개의 가장 합이 작은 pair 를 반환해야한다.

풀이

1️⃣ PriortyQueue 와 일종의 그리디 법을 활용

💡 참고한 Idea

이번 문제는, 아무리 열심히 떠올려봐도 시간 초과 & 메모리 초과 때문에 해결하지 못하고 다른 사람의 아이디어를 참고하여 해결하였다.

핵심 아이디어는 다음과 같다.

먼저, 두 원소의 합을 기준으로 오름차순으로 정렬하는 우선순위 큐를 준비한다.

nums1 배열의 인덱스는 0 으로 고정한 채, nums2 배열을 순회하며 합을 구하고 우선순위 큐에 입력한다.

우선순위 큐에는 [ nums1의 값, nums2 의 값, nums1의 인덱스, nums1의 값 +nums2의 값 ] 으로 입력해준다.

이제 우선순위 큐에서 poll()을 하는데 여기서 poll() 한 값은 큐에서 두 원소의 합이 가장 작은 값이 추출될 것이다.

그렇게 poll() 한 값에서 nums1의 인덱스를 1 추가해주고 (nums1 배열의 길이를 초과하지 않는 한)

우선순위 큐에 다시 넣어준다.

이 값을 다시 넣어주는 이유는, 이 값은 제거당한 값과 가장 비슷한 위치의 다른 쌍이 될 수 있고 이 합이 이미 들어간 다른 쌍보다 적은 합이 나올 수 있기 때문이다.

말로 설명하니 약간 이해하기 힘들 수 있는데, 코드를 보면 이해가 잘 될 것이다.

import java.util.*;

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> ans = new ArrayList<>();
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[3] - b[3]);

        for (int i = 0; i < Math.min(k, nums2.length); i++) {
            pq.offer(new int[]{nums1[0], nums2[i], 0, nums1[0] + nums2[i]});
        }

        while (!pq.isEmpty() && k-- > 0) {
            int[] poll = pq.poll();
            int num1 = poll[0], num2 = poll[1];
            ans.add(List.of(num1, num2));

            int idx1 = poll[2];
            if (idx1 + 1 < nums1.length) {
                // num2와 제거한 num1 뒤에 있는 수의 조합을 추가
                pq.offer(new int[]{nums1[idx1 + 1], num2, idx1 + 1, nums1[idx1 + 1] + num2});
            }
        }


        return ans;
    }
}