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 (u1,v1),(u2,v2) ...(uk,vk) 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]

Hint

Train of Thought

Several solutions from naive to more elaborate. I found it helpful to visualize the input as an m×n matrix of sums, for example for nums1=[1,7,11], and nums2=[2,4,6]:

  2   4   6

+------------ 1 | 3 5 7 7 | 9 11 13 11 | 13 15 17

Solution 1: Brute Force (accepted in 560 ms) Just produce all pairs, sort them by sum, and return the first k. Solution 2: Clean Brute Force (accepted in 532 ms)

The above produces tuples and while the judge doesn't care, it's cleaner to make them lists as requested:

Solution 4: Efficient (accepted in 112 ms)

The brute force solutions computed the whole matrix (see visualization above). This solution doesn't. It turns each row into a generator of triples [u+v, u, v], only computing the next when asked for one. And then merges these generators with a heap. Takes O(m + k*log(m)) time and O(m) extra space.

Solution 5: More efficient (accepted in 104 ms)

The previous solution right away considered (the first pair of) all matrix rows (see visualization above). This one doesn't. It starts off only with the very first pair at the top-left corner of the matrix, and expands from there as needed. Whenever a pair is chosen into the output result, the next pair in the row gets added to the priority queue of current options. Also, if the chosen pair is the first one in its row, then the first pair in the next row is added to the queue.

Code

public class Solution {

    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> list = new ArrayList();
        int m = nums1.length;
        int n = nums2.length;
        if(m == 0 || n == 0 || k < 1) return list;

        PriorityQueue<Data> heap = new PriorityQueue<Data>((a, b) -> (a.val - b.val));

        heap.offer(new Data(0, 0, nums1[0] + nums2[0]));

        while(!heap.isEmpty() && k > 0){
            Data d = heap.poll();

            int[] pair = {nums1[d.i], nums2[d.j]};
            list.add(pair);
            k--;

            //add the next column item
            if(d.j < n - 1){
                heap.offer(new Data(d.i, d.j + 1, nums1[d.i] + nums2[d.j + 1]));
            }
            // always store the next row (smallest) 
            // 规定只有第一列可以向下走, 其他列只能往右走, 这样避免重复计算
            if(d.j == 0 && d.i < m - 1){
                heap.offer(new Data(d.i + 1, 0, nums1[d.i + 1] + nums2[0]));
            }
        }
        return list;
    }

        class Data{
            int i; 
            int j;
            int val;
            public Data(int i, int j, int val){
                this.i = i;
                this.j = j;
                this.val = val;
            }
        }
}

Complexity