373. Find K Pairs with Smallest Sums

Description

Given two sorted arrays nums1 and nums2 of sizes m and n respectively, return an array of all the pairs [nums1[i], nums2[j]] such that i < m and j < n and the sum of nums1[i] + nums2[j] is minimized. The number of pairs to return is k.

Intuition

To find the k pairs with the smallest sums efficiently, we can use a min-heap (priority queue). By leveraging the properties of the min-heap, we can always extract the smallest pair and keep track of the next possible pair to consider from the sorted arrays.

Approach

1. Initialize a min-heap where each element is an array containing three integers: the sum of the pair, the index in nums1, and the index in nums2.
2. Push the first pair from each element in nums1 paired with the first element in nums2 into the min-heap.
3. Extract the smallest pair from the heap, add it to the result list, and then push the next pair that includes the same element from nums1 but the next element from nums2.
4. Repeat until we have extracted k pairs or the heap is empty.

Complexity

Time Complexity:

O(k log k), because we perform k operations of extracting from and adding to the heap, each taking O(log k) time.

Space Complexity:

O(k), because the heap will hold up to k elements.

Code

C++

class Solution {
public:
    vector> kSmallestPairs(vector& nums1, vector& nums2, int k) {
        vector> result;
        auto compare = [&nums1, &nums2](pair& a, pair& b) {
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        };
        priority_queue, vector>, decltype(compare)> minHeap(compare);

        for (int i = 0; i < nums1.size() && i < k; ++i) {
            minHeap.emplace(i, 0);
        }

        while (k > 0 && !minHeap.empty()) {
            auto [i, j] = minHeap.top();
            minHeap.pop();
            result.push_back({nums1[i], nums2[j]});

            if (j + 1 < nums2.size()) {
                minHeap.emplace(i, j + 1);
            }

            --k;
        }

        return result;
    }
};

Python

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        import heapq
        result = []
        min_heap = []

        for i in range(min(len(nums1), k)):
            heapq.heappush(min_heap, (nums1[i] + nums2[0], i, 0))

        while k > 0 and min_heap:
            _, i, j = heapq.heappop(min_heap)
            result.append([nums1[i], nums2[j]])
            if j + 1 < len(nums2):
                heapq.heappush(min_heap, (nums1[i] + nums2[j + 1], i, j + 1))
            k -= 1

        return result

Java

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

        for (int i = 0; i < nums1.length && i < k; ++i) {
            minHeap.offer(new int[] { i, 0 });
        }

        while (k > 0 && !minHeap.isEmpty()) {
            int[] top = minHeap.poll();
            int i = top[0], j = top[1];
            result.add(Arrays.asList(nums1[i], nums2[j]));
            if (j + 1 < nums2.length) {
                minHeap.offer(new int[] { i, j + 1 });
            }
            --k;
        }

        return result;
    }
}

JavaScript

/**
 * @param {number[]} nums1
 * @param {number[]} nums2
 * @param {number} k
 * @return {number[][]}
 */
var kSmallestPairs = function(nums1, nums2, k) {
    const result = [];
    const heap = new MinPriorityQueue({ compare: (a, b) => a[0] - b[0] });
    
    for (let i = 0; i < nums1.length; i++) {
        heap.enqueue([nums1[i] + nums2[0], 0]);
    }
    
    while (k > 0 && !heap.isEmpty()) {
        const [nums_sum, index] = heap.dequeue();
        result.push([nums_sum - nums2[index], nums2[index]]);
        
        if (index + 1 < nums2.length) {
            heap.enqueue([nums_sum - nums2[index] + nums2[index + 1], index + 1]);
        }
        k--;
    }
    
    return result;
};