373. Find K Pairs with Smallest Sums (Medium)

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]

Solution: Brute Force

Suppose $$N = min(k, n1*n2)$$

Time complexity: $$O(N^2)$$ 502 ms

这道题给了我们两个数组，让我们从每个数组中任意取出一个数字来组成不同的数字对，返回前K个和最小的数字对。那么这道题有多种解法，我们首先来看brute force的解法，这种方法我们从0循环到数组的个数和k之间的较小值，这样做的好处是如果k远小于数组个数时，我们不需要计算所有的数字对，而是最多计算k*k个数字对，然后将其都保存在res里，这时候我们给res排序，用我们自定义的比较器，就是和的比较，然后把比k多出的数字对删掉即可，参见代码如下：

v1: QuickSort with vector

vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
    vector<pair<int, int>> res;
    for (int i = 0; i < min((int)nums1.size(),k); ++i) {
        for (int j = 0; j < min((int)nums2.size(),k); ++j) {
            res.push_back({nums1[i],nums2[j]});
        }
    }

    sort(res.begin(), res.end(), [](pair<int, int> &a, pair<int, int> &b){return a.first+a.second < b.first+b.second;});

    if (res.size() > k) return vector<pair<int, int>>(res.begin(), res.begin()+k);
    return res;
}

v2: priority_queue

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> res;
        priority_queue<pair<int, int>, vector<pair<int, int>>, cmp> q;
        for (int i = 0; i < min((int)nums1.size(), k); ++i) {
            for (int j = 0; j < min((int)nums2.size(), k); ++j) {
                if (q.size() < k) {
                    q.push({nums1[i], nums2[j]});
                } else if (nums1[i] + nums2[j] < q.top().first + q.top().second) {
                    q.push({nums1[i], nums2[j]}); q.pop();
                }
            }
        }
        while (!q.empty()) {
            res.push_back(q.top()); q.pop();
        }
        return res;
    }
    struct cmp {
        bool operator() (pair<int, int> &a, pair<int, int> &b) {
            return a.first + a.second < b.first + b.second;
        }
    };
};

Solution 2: Heap 13ms

$$N = min(k, n1*n2)$$

Space Complexity: $$O(N)$$

Time Complexity: $$O(NlogN)$$

Suppose you are at pair 0,0 (index 0 and index 0, not value), which is the current minimum.
Then you know the only two possible followers (immediate larger ones) are 0,1 and 1,0. Any other indices, say 1,1, 1,2, 3,3 have to be larger right?
So every time you get a current minimum i,j , you want to push i+1,j and i,j+1 into heap. You don't need to worry about others yet.

The problem here is, if you are at 2,3, you will push 3,3 and 2,4; then later, you are at 3,2. Then you will push 3,3, 4,2. so you pushed 3,3 twice.

But it is easy to realize, if you are at 2,3, and you haven't seen 3,2 yet (meaning 3,2 is either still in the queue but not popped yet or not even pushed into queue), you don't need to worry about 3,3 at this moment, because 3,2 is guaranteed to be no greater than 3,3. So you can wait till you see 3,2.

This basically means every time you see i,j you just need to push i+1, j into queue. i, j+1 can be pushed into queue later when you see i - 1, j + 1. The only exception to this is, when i == 0, there is no i-1, j+1 anymore, so you want to push both i+1, j and i, j+1 when i == 0.

vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
    vector<pair<int, int>> res;
    if (nums1.empty() || nums2.empty() || k <= 0) return res;
    auto cmp = [&nums1, &nums2](pair<int, int> &a, pair<int, int> &b)
        {return nums1[a.first]+nums2[a.second] > nums1[b.first]+nums2[b.second];};
    priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> heap(cmp); // min heap

    heap.push({0,0});
    int n1 = nums1.size(), n2 = nums2.size();
    while (k-- && heap.size()) {
        auto idx = heap.top(); heap.pop(); // pop the minimum
        res.push_back({nums1[idx.first], nums2[idx.second]});
        if (idx.first + 1 < n1) heap.push({idx.first+1, idx.second});
        if (idx.first == 0 && idx.second + 1 < n2)
            heap.push({idx.first, idx.second+1});
    }
    return res;
}