Find K Pairs with Smallest Sums

Problem

You are given two integer arrays nums1 and nums2 sorted in ascending 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.

 

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]
Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]
Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

 

Constraints:

• 1 <= nums1.length, nums2.length <= 105
• -109 <= nums1[i], nums2[i] <= 109
• nums1 and nums2 both are sorted in ascending order.
• 1 <= k <= 1000

Solution

/**
 * @param {number[]} nums1
 * @param {number[]} nums2
 * @param {number} k
 * @return {number[][]}
 */
var kSmallestPairs = function(nums1, nums2, k) {
    const minHeap = new MinPriorityQueue({ priority: tuple => tuple.v1 + tuple.v2 });
    const res = [];

    // push front of 'list' into heap
    for (let i = 0; i < k && i < Math.min(nums2.length, k); i++) {
        minHeap.enqueue({ i1: 0, i2: i, v1: nums1[0], v2: nums2[i] });
    }

    // find k smallest pairs
    while (res.length < k && !minHeap.isEmpty()) {
        const t = minHeap.dequeue().element;
        res.push([t.v1, t.v2]);
        if (t.i1 + 1 < nums1.length) {
            minHeap.enqueue({ i1: t.i1 + 1, i2: t.i2, v1: nums1[t.i1 + 1], v2: nums2[t.i2] });
        }
    }
    return res;
};

The solution is extremely similar to Merge k Sorted Lists. Observe that the k smallest pairs must be in the set obtained from the cardinal product of the first k elements in nums1 with the first k elements in nums2. There are a total of k² pairs, which we can use to simulate k sorted lists. More specifically, the first list is the first element in nums1 paired with the first k elements in nums2, the second list is the second element in nums1 pairs with the first k elements in nums2, and so on (each list is sorted in ascending order with respect to the sum).

To optimize this process, we can use a minHeap that contains a max of k pairs of elements at any given time. Initially, minHeap will contain the min(nums2.length, k) pairs (v1, v2) which are obtained from the cardinal product of the two sets nums1[0:1] and nums2[0:k + 1].