[LeetCode] Find K Pairs with Smallest Sums 找和最小的K对数字

 

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]

 

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

 

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

 

解法一：

class Solution {
public:
    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) res.erase(res.begin() + k, res.end());
        return res;
    }
};

 

我们也可以使用multimap来做，思路是我们将数组对之和作为key存入multimap中，利用其自动排序的机制，这样我们就可以省去sort的步骤，最后把前k个存入res中即可：

 

解法二：

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> res;
        multimap<int, pair<int, int>> m;
        for (int i = 0; i < min((int)nums1.size(), k); ++i) {
            for (int j = 0; j < min((int)nums2.size(), k); ++j) {
                m.insert({nums1[i] + nums2[j], {nums1[i], nums2[j]}});
            }
        }
        for (auto it = m.begin(); it != m.end(); ++it) {
            res.push_back(it->second);
            if (--k <= 0) return res;
        }
        return res;
    }
};

 

下面这种方式用了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;
        }
    };
};

 

下面这种方法比较另类，我们遍历nums1数组，对于nums1数组中的每一个数字，我们并不需要遍历nums2中所有的数字，实际上，对于nums1中的数字，我们只需要记录nums2中下一个可能组成数字对的坐标：

 

解法四：

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> res;
        int size = min(k, int(nums1.size() * nums2.size()));
        vector<int> idx(nums1.size(), 0);
        for (int i = 0; i < size; ++i) {
            int t = 0, sum = INT_MAX;
            for (int j = 0; j < nums1.size(); ++j) {
                if (idx[j] < nums2.size() && sum >= nums1[j] + nums2[idx[j]]) {
                    t = j;
                    sum = nums1[j] + nums2[idx[j]];
                }
            }
            res.push_back({nums1[t], nums2[idx[t]]});
            ++idx[t];
        }
        return res;
    }
};

 

参考资料：

https://discuss.leetcode.com/topic/50429/c-solution

https://discuss.leetcode.com/topic/50459/c-idea-of-using-multimap

https://discuss.leetcode.com/topic/50422/naive-accepted-solution-c

https://discuss.leetcode.com/topic/50421/c-priority_queue-solution

 

LeetCode All in One 题目讲解汇总(持续更新中...)