[Leetcode] Find K Pairs with Smallest Sums

Perhaps the most counter-intuitive heap problem I’ve done so far.

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.

Return the k pairs (u1, v1), (u2, v2), ..., (uk, vk) with the smallest sums.

Examples

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]

Solution

The essence is keeping the least number of pairs in a container that guarantees the smallest pair of all can be chosen from the container. Heap is used as container because only smallest is needed rather than sort all container.

Set nums1=short list & nums2=long list. Consider how to init the heap. First, if k<len(nums1), use nums2[0] with top k elm of nums1 to form k pairs and push them into heap. Why only k? For the k+1th elm, the smallest possible pair is (k+1th, 0), and it’s only required if top K pairs have been used. However, once top K pairs have been used, the condition to return is satisfied and k+1th will not be reached.

If k≥len(nums1), then nums2[0] with all elms of nums1 are needed to. Why can’t use fewer? If n elms are used where n<min(k, len(nums1)), when n pairs are used, we will find that the condition to return is not satisfied. As we need to push new elm in, the problem arises that (n+1th, 0) is not guaranteed to be smaller than (0~nth, >0). Hence, newly pushed pairs may not be able to contain the smallest pair.

After init the heap, we pop the smallest pair, and push new pair in. How to choose new? For each popped elm(i, j), push (i, j+1). This is tricky because (i, j+1) may not be the smallest of rest. However, say another pair(i+x, j)<(i, j+1), it does not matter as long as heap contains (<i+x, j). That is to say, though there may be other smaller pairs, heap surely contains smaller pair of these smaller pairs, so the smallest pair with same i of the popped is sufficient.

• code

ref: https://blog.csdn.net/fuxuemingzhu/article/details/79564362