[ 오늘의 코테 연습장 ] [ LeetCode ] - 373. Find K Pairs with Smallest Sums

📑 문제

두 개의 정수 배열과 정수 k가 주어졌을 때, 두 배열에서 각각 하나씩 선택하여 만들 수 있는 쌍(pair) 중 합이 가장 작은 k개의 쌍을 찾는 문제입니다.

두 개의 입력 리스트 nums1, nums2와 정수 k가 주어집니다.
각 리스트에서 하나의 원소를 선택하여 쌍을 만들 수 있습니다.
이러한 쌍 중에서 합이 가장 작은 k개의 쌍을 찾아야 합니다.
결과를 반환할 때, 각 쌍은 list로 구성되며, 전체 결과는 list of lists 형태로 반환해야 합니다.

📑 문제 접근 방법

처음에는 뭔가 쉽게 풀수 있는 문제 같았다.

근데 각 배열의 요소를 조합해서 만든 쌍 중에 합이 가장 작은 쌍을 찾아야하니깐

두 배열에서 만들 수 있는 쌍의 조합을 만들고 , 그중에서 가장 작은 쌍을 찾을려고했는데 이렇게 구현하게 된다면 시간복잡도가 너무 커져서 time out이 날 것 같아 다른 방법을 찾았다.

그래서 나는 Heap을 이용해 문제를 해결할려고 했다.

📑 Heap, 우선순위 큐 사용 이유

문제에서는 가장 작은 합을 갖는 쌍을 찾아야 하는데, 이러한 연산에 Heap이 매우 효율적이기 때문에 사용하려고 했다.

우선순위 큐는 내부적으로 Heap을 사용해서 원소들을 자동으로 정렬할 수 있기 때문에, 모든 가능한 쌍을 조합하고, 정렬할 필요없이 우선순위 큐를 이용해 K번의 연산만으로 타임 아웃없이 문제를 해결할 수 있어서 선택하게 되었다.

📑 code

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new ArrayList<>();
        
        if (nums1.length == 0 || nums2.length == 0 || k == 0) {
            return result;
        }
        
        // 우선순위 큐 ( 오름차순 )
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(
            (a, b) -> a[0] + a[1] - b[0] - b[1]
        );
        int sizeOfNum1=  Math.min(nums1.length, k);

        for (int i = 0; i <sizeOfNum1; i++) {
            minHeap.offer(new int[]{nums1[i], nums2[0], 0});
        }
        
        while(k-- > 0 && !minHeap.isEmpty()) {
            int[] cur = minHeap.poll();
            result.add(List.of(cur[0], cur[1]));
            
            if(cur[2] == nums2.length - 1) continue;
            
            // 현재 쌍에서 nums1의 값과 인덱스는 그대로 유지하고, nums2의 인덱스만 1 증가시킨다.
           int nextNums2Index = cur[2] + 1;

            // 같은 nums1 값과 다음 nums2 값을 짝지어서 minHeap에 추가한다.
            minHeap.offer(new int[]{cur[0], nums2[nextNums2Index], nextNums2Index});
        }
        
        return result;
    }
}