利用堆,维护当前最小值。
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
auto cmp = [](pair<int, int> x, pair<int, int> y){
return x.first + x.second < y.first + y.second;};
priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> heap(cmp);
vector<vector<int>> result;
for(int i = 0; i < nums1.size() && i < k; i++){
for(int j=0; j < nums2.size() && j < k; j++){
auto tmp = make_pair(nums1[i], nums2[j]);
if(heap.size() < k){
heap.push(tmp);
} else {
if(!cmp(heap.top(), tmp)){
heap.pop();
heap.push(tmp);
}
}
}
}
while(!heap.empty()){
result.push_back({heap.top().first, heap.top().second});
heap.pop();
}
return result;
}
一个别人的写法,学习一下代码风格
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
if (nums1.empty() || nums2.empty())
return vector< vector<int> >();
using Pair = pair<int, int>;
auto cmp_func = [&](Pair a, Pair b){
return (nums1[a.first] + nums2[a.second]) >
(nums1[b.first] + nums2[b.second]);
};
using Pri = priority_queue<Pair, vector<Pair>, decltype(cmp_func)>;
Pri minheap(cmp_func);
vector< vector<int> > result;
for (int i = 0; i < nums1.size() && i < k; i++)
{
minheap.push({i, 0});
}
while (result.size() < k && !minheap.empty())
{
auto min = minheap.top();
minheap.pop();
result.push_back({nums1[min.first], nums2[min.second]});
if (min.second < nums2.size() - 1)
{
minheap.push({min.first, min.second + 1});
}
}
return result;
}