373. Find K Pairs with Smallest Sums

Description

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.

Find the k pairs (u₁,v₁),(u₂,v₂) ...(u_k,v_k) with the smallest sums.

Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6], k = 3

Return: [1,2],[1,4],[1,6]

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:

Given nums1 = [1,1,2], nums2 = [1,2,3], k = 2

Return: [1,1],[1,1]

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:

Given nums1 = [1,2], nums2 = [3], k = 3

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

Credits:
Special thanks to @elmirap and @StefanPochmann for adding this problem and creating all test cases.

Solution

MinHeap, time O(k * log k), space O(k)

跟“378. Kth Smallest Element in a Sorted Matrix”的解法完全一样。

这种求kth的题目一定要反应到Heap上...这道题用m * n的heap也能做，但是更好的解法是用大小为k的heap做。因为这道题里面sum是有规律的，可以保证第k次访问queue时，kth元素一定在queue中且被取出就行了。

注意heap的大小。预添加元素到heap的时候，如果k >= nums2.length，则把nums2全部offer进去；但如果k < nums2.length，则只offer nums2的前k个元素进去即可，因为后面的元素绝对不会被用到了。利用这点可以尽量缩小heap大小，时间和空间都有节省。

class Solution {
    public List kSmallestPairs(int[] nums1, int[] nums2, int k) {
        if (nums1 == null || nums2 == null 
            || nums1.length < 1 || nums2.length < 1 || k < 1) {
            return Collections.EMPTY_LIST;
        }
        
        PriorityQueue queue 
            = new PriorityQueue<>((a, b) -> (a.val - b.val));
        // offer k pairs formed by nums1[0] and nums2[j] to queue
        for (int j = 0; j < Math.min(nums2.length, k); ++j) {
            queue.add(new Tuple(0, j, nums1[0] + nums2[j]));
        }
        
        List kSmallestPairs = new ArrayList<>();
        while (k-- > 0 && !queue.isEmpty()) {
            Tuple t = queue.poll();
            kSmallestPairs.add(new int[]{nums1[t.i], nums2[t.j]});
            
            if (t.i == nums1.length - 1) {
                continue;
            }
            queue.offer(new Tuple(t.i + 1, t.j, nums1[t.i + 1] + nums2[t.j]));
        }
        
        return kSmallestPairs;
    }
    
    class Tuple {
        int i;
        int j;
        int val;
        
        public Tuple(int i, int j, int val) {
            this.i = i;
            this.j = j;
            this.val = val;
        }
    }
}