373 Find K Pairs with Smallest Sums

373. Find K Pairs with Smallest Sums

1. 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.

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]

2. Implementation

(1) Heap

class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> res = new ArrayList<>();

        if (nums1 == null || nums1.length == 0 || nums2 == null || nums2.length == 0) {
            return res;
        }

        PriorityQueue<PairSum> minHeap = new PriorityQueue<>();

        int m = nums1.length, n = nums2.length;

        for (int j = 0; j < n; j++) {
            minHeap.add(new PairSum(0, j, nums1[0] + nums2[j]));
        }

        for (int i = 0; i < Math.min(k, m * n); i++) {
            PairSum curSum = minHeap.remove();
            int index1 = curSum.index1;
            int index2 = curSum.index2;

            res.add(new int[] {nums1[index1], nums2[index2]});

            if (index1 < m - 1) {
                minHeap.add(new PairSum(index1 + 1, index2, nums1[index1 + 1] + nums2[index2]));
            }
        }
        return res;
    }

    class PairSum implements Comparable<PairSum> {
        int index1;
        int index2;
        int sum;

        public PairSum(int index1, int index2, int sum) {
            this.index1 = index1;
            this.index2 = index2;
            this.sum = sum;
        }

        public int compareTo(PairSum that) {
            return this.sum - that.sum;
        }
    }
}

3. Time & Space Complexity

Heap: 时间复杂度O(k * k *logk), 空间复杂度O(k)