LeetCode 061. 剑指 Offer II 061. 和最小的 k 个数对(多路归并)
原题链接
中等
作者:
我要出去乱说
,
2022-04-03 10:49:41
,
所有人可见
,
阅读 240
1、思路
- 建一个小跟堆,堆中存放
pair<int,int>
,pair
中分别是两个数组元素的下标,自定义比较函数,使得堆顶pair
中所代表的数组和最小;
- 初始化堆中元素:将第一个数组的每个元素与第二个数组中的第一个元素配对,分别插入堆中;
- 循环弹出堆顶元素,每次弹出后,就把第一数组与第二数组后一元素的下标插入堆中,堆会自动排序。
2、代码
class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
int n = nums1.size(), m = nums2.size();
vector<vector<int>> res;
//因为在lamba表达式中默认不能够使用任何外部变量,只有写在[]中的外部变量才能够使用
//[=]表示按值捕获所有外部变量,[&]表示按引用捕获所有外部变量
auto cmp = [&nums1, &nums2] (const pair<int, int>& a, const pair<int, int>& b) {
return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
};
//此处decltype是用来定义函数类型的
priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> q(cmp);
//只需要求最小的k个数对,又因为数组是升序,故不需要考虑下标在k之后的元素,所以用min(k, n)
for (int i = 0; i < min(k, n); ++ i)
{
q.push(make_pair(i, 0)); //初始化堆
}
while (k -- && !q.empty())
{
auto [x, y] = q.top(); //多元素赋值必须用auto
q.pop();
res.push_back({nums1[x], nums2[y]});
if (y < m - 1)
{
q.push(make_pair(x, y + 1));
}
}
return res;
}
};