373. Find K Pairs with Smallest Sums

 

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.

 

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]

 

Constraints:

  • 1 <= nums1.length, nums2.length <= 104
  • -109 <= nums1[i], nums2[i] <= 109
  • nums1 and nums2 both are sorted in ascending order.
  • 1 <= k <= 1000

 

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 中下一个可能组成数字对的坐标,这里使用一个 idx 数组,其中 idx[i] 表示的数字是 nums1[i] 将从 nums2 数组上开始寻找的位置,因为 {nums1[i], nums2[i - 1]} 已经被加入到了结果 res 中,这种方法其实也是一种地毯式搜索,但是并不需要遍历完所有的组合,因为有 idx 数组来进行剪枝。suppose 需要进行k次循环,但是题目中没有说一定能取出k对数字,而能取出的对儿数跟数组 nums1 和 nums2 的长度有关,最多能取出二者的长度之积的对儿数,所以取其跟k之间的较小值为循环次数。我们定义 idx 数组,长度为 nums1 的长度,初始化均为0。下面开始循环,在每次循环中,新建变量 cur,记录从 nums1 中取数的位置,初始化为0,使用变量 sum 来记录一个当前最小的两数之和,初始化为正无穷。然后开始遍历数组 nums1,更新 sum 的条件有两个,第一个是 idx[i] 上的数要小于 nums2 的长度,因为其是在 nums2 开始寻找的位置,当然不能越界,第二个条件的候选的两个数组 nums1[i] 和 nums2[idx[i]] 之和小于等于 sum。同时满足这两个条件就可以更新 sum了,同时更新 cur 为i,表示当前从 nums1 取出数字的位置。当遍历 nums1 的 for 循环结束后,此时 cur 的位置就是要从 nums1 取出的数字的位置,根据 idx[cur] 从 nums2 中取出对应的数组,形成数对儿存入结果 res 中,然后 idx[cur] 自增1,因为当前位置的数字已经用过了,下次遍历直接从后面一个数字开始吧,这是本解法的设计精髓所在,一定要弄清楚 idx 数组的意义,参见代码如下:

 

解法四:

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 t = 0; t < size; ++t) {
            int cur = 0, sum = INT_MAX;
            for (int i = 0; i < nums1.size(); ++i) {
                if (idx[i] < nums2.size() && sum >= nums1[i] + nums2[idx[i]]) {
                    cur = i;
                    sum = nums1[i] + nums2[idx[i]];
                }
            }
            res.push_back({nums1[cur], nums2[idx[cur]]});
            ++idx[cur];
        }
        return res;
    }
};

 

Github 同步地址:

https://github.com/grandyang/leetcode/issues/373

 

参考资料:

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/discuss/84655/c-solution

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/discuss/84653/c-idea-of-using-multimap

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/discuss/84551/simple-Java-O(KlogK)-solution-with-explanation

 

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


转载请注明来源于 Grandyang 的博客 (grandyang.com),欢迎对文章中的引用来源进行考证,欢迎指出任何有错误或不够清晰的表达。可以在下面评论区评论,也可以邮件至 grandyang@qq.com

💰


微信打赏


Venmo 打赏

(欢迎加入博主的知识星球,博主将及时答疑解惑,并分享刷题经验与总结,试运营期间前五十位可享受半价优惠~)

×

Help us with donation