前言
题目描述
给定两个以升序排列的整形数组 nums1 和 nums2, 以及一个整数 k。
定义一对值 (u,v),其中第一个元素来自 nums1,第二个元素来自 nums2。
找到和最小的 k 对数字 (u1,v1), (u2,v2) … (uk,vk),按从小到大的顺序输出它们的和。
解题思路
其实我们要有一个先验知识,就是,遇见K问题的时候,要想到优先队列,包括大顶堆和小顶堆;
回顾:
像我们上一题中“前K个高频元素”,我们创建小顶堆,将前K个高频的元素都入堆,当新元素来临时,只要比较堆顶和新元素的大小即可,如果新元素比堆顶大,那么就把堆顶弹出,新元素入堆,如果新元素小,则说明堆中的元素频度都比新元素大,则不处理,遍历下一个值就行。
那么这道题类似:
和最小,那就创建大顶堆,因为只要新元素比堆顶大,那么就比堆中所有的元素大,此时就不用处理,遍历下一个就行,但是如果新元素比堆顶元素小,就需要将堆顶元素弹出,然后将新元素入堆。依次我们设定堆中只能有k个元素,最终堆中的元素就是我们的答案。
代码
#include<iostream>
#include<unordered_map>
#include<vector>
#include<queue>
#include<algorithm>
#include<array>
using namespace std;
class Solution{
public:
vector<vector<int>> KSmallestPairs(vector<int> &nums1,vector<int> &nums2,int k){
//创建一个大顶堆,由于优先队列默认的就是大顶堆,我们就不必像上一题那样再重新写比较函数
priority_queue<pair<int,vector<int>>> q;
for(auto a:nums1){
for(auto b:nums2){
if(q.size() < k){
q.push({a+b,{a,b}});
}
else if(a+b < q.top().first){
q.pop();
q.push({a+b,{a,b}});
}
}
}
vector<vector<int>> res;
while(!q.empty()){
res.push_back(q.top().second);
q.pop();
}
//逆置
reverse(res.begin(),res.end());
return res;
}
};
int main(){
int n;
cin>>n;
int data;
vector<int> nums1;
for(int i=0;i<n;i++){
cin>>data;
nums1.push_back(data);
}
int m;
cin>>m;
int data2;
vector<int> nums2;
for(int i=0;i<m;i++){
cin>>data2;
nums2.push_back(data2);
}
int k;
cin>>k;
//按从小到大的顺序输出k对数字的和
vector<vector<int>> res=Solution().KSmallestPairs(nums1,nums2,k);
for(int i=0;i<res.size();i++){
if(i>0){
cout<<" ";
}
cout<<res[i][0]+res[i][1];
}
return 0;
}