Find K Pairs with Smallest Sums

문제

You are given two integer arrays nums1 and nums2 sorted in non-decreasing order and an integer k.

Define a pair (u, v) which consists of one element from the first array and one element from the second array.

Return the k pairs (u1, v1), (u2, v2), ..., (uk, vk) with the smallest sums.

풀이

nums1, nums2 q배열에서 각 원소를 더해 작은 값이 되는 수를 k개만 반환하는 문제

priority queue를 이용해 풀이할 수 있다.

• 먼저 pq에 {nums1의 원소 + nums2[0], 0} 형태로 추가
• pq가 빌 때까지 값을 꺼내오면서 answer list에 추가
  • 도중에 nums2에 값이 더 남아있을 경우 pq에 추가

수도 코드

결과값을 담을 answer list 생성
heap 생성

for nums1의 원소 num1:
	pq에 {num1 + nums2[0], 0} 추가

while k가 0보다 크고 pq가 비지 않았다면:
	sum, pos = pq.poll()

    answer에 {sum - nums2[pos], nums2[pos]} 추가
    
    if nums2를 더 탐색할 수 있다면:
    	pq에 {sum - nums2[pos] + nums2[pos + 1], pos + 1} 추가
    
    k--
    
return answer

Solution(Runtime: 28ms)

import java.util.*;


class Solution {

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

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

        while (k > 0 && !pq.isEmpty()) {
            int[] now = pq.poll();
            int sum = now[0], pos = now[1];

            answer.add(List.of(sum - nums2[pos], nums2[pos]));

            if (pos < nums2.length - 1) {
                pq.add(new int[] {sum - nums2[pos] + nums2[pos + 1], pos + 1});
            }
            k--;
        }

        return answer;
    }
}

첫번째 단계에서 nums1, nums2의 모든 원소를 더한다면 시간초과가 발생한다. 당연하게도 두 배열의 최대 길이는 $10^5$이기 때문이다.

따라서 저 단계에서 어떤 값을 넣어야할지가 중요해진다. 따라서 nums2의 모든 원소를 더하지 않고 nums2는 고정하고 nums1의 값만 늘려서 추가했다. 마지막값은 nums2의 인덱스이다.

2번째 단계에서는 pq의 저장된 최솟값을 가져와 answer에 {nums1의 원소, nums2의 원소} 형태로 저장해 주었다.

그리고 nums2의 탐색을 담당하는 pos 변수가 nums2를 더 탐색할 수 있는 경우에는 pq에 다음 값도 추가할 수 있도록 하였다.