373. Find K Pairs with Smallest Sums

문제

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

풀이

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> heap = new PriorityQueue<>((a,b) -> {
            return (nums1[a[0]] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]]);
        });
        List<List<Integer>> list = new ArrayList<>();


        // 초기 힙 구성
        for(int i = 0; i < nums2.length; i++) {
            heap.add(new int[]{i,0});
        }

        // k번만큼만 반복
        while(k-- > 0 && !heap.isEmpty()) {
            int[] curr = heap.poll();
            int i = curr[0];
            int j = curr[1];
            list.add(Arrays.asList(nums1[i], nums2[j]));

            if(j + 1 < nums2.length) {
                heap.add(new int[]{i, j+1});
            }
        }

        return list;
    }
}

• 핵심은 nums1과 nums2를 돌면서 합을 전부 구하는 방식처럼 O(N^2)으로 풀지 않는 것.
• 초기 힙에 (i, 0)을 전부 넣어준 뒤,
  (i, j)가 나왔으면, 그 다음 가능한 작은 조합은 (i, j+1)일 수밖에 없음