LeetCode 373. Find K Pairs with Smallest Sums

Question

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 <= 104

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

Solution

This problem is a disguised LeetCode 378. Kth Smallest Element in a Sorted Matrix. We can treat combinations of nums1 and nums2 as a (len1,len2) matrix. Each element in this matrix is the sum of a pair. Also, this matrix is sorted, which means that every row and column is sorted in ascending order.

static class element {
    // nums1[x] + nums2[y]
    int val;
    int x;
    int y;

    public element(int _val, int _x, int _y) {
        this.val = _val;
        this.x = _x;
        this.y = _y;
    }
}

// min heap, Time Complexity O(k*log(k))
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    int len1 = nums1.length;
    int len2 = nums2.length;
    if (len1 == 0 || len2 == 0 || k == 0) {
        return new ArrayList<>();
    }

    PriorityQueue<element> minHeap = new PriorityQueue<>(Math.min(k, len1), new Comparator<element>() {
        @Override
        public int compare(element o1, element o2) {
            return o1.val - o2.val;
        }
    });
    for (int i = 0; i < Math.min(k, len1); i++) {
        minHeap.add(new element(nums1[i] + nums2[0], i, 0));
    }

    List<List<Integer>> result = new ArrayList<>();
    // treat combinations of nums1 and nums2 as a (len1,len2) matrix
    for (int count = 0; count < k && !minHeap.isEmpty(); count++) {
        element e = minHeap.poll();
        result.add(new ArrayList<Integer>(Arrays.asList(nums1[e.x], nums2[e.y])));
        if (e.y < len2 - 1) {
            minHeap.add(new element(nums1[e.x] + nums2[e.y + 1], e.x, e.y + 1));
        }
    }
    return result;
}