使用最笨的办法，直接用一个优先级队列实现，从最开始往优先级队列中填充元素，直到满了以后，就把最大的那个元素出队，这样最终就可以得到一个最小和的 pair 的集合，这样的时间效率是 O(n^2)。显然这样有些慢。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
| class Solution {
private:
struct cmp{
bool operator()(pair<int,int>& x,pair<int,int>& y){
return x.first+x.second<y.first+y.second;
}
};
public:
vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
vector<pair<int,int>> result;
if(k<=0) return result;
priority_queue<pair<int,int>,vector<pair<int,int>>,cmp> pq;
for(int i=0;i<nums1.size();i++){
for(int j=0;j<nums2.size();j++){
pq.push(make_pair(nums1[i],nums2[j]));
if(pq.size()>k) pq.pop();
}
}
result.resize(pq.size());
for(int i=pq.size()-1;i>=0;i--){
result[i]=pq.top();
pq.pop();
}
return result;
}
};
|
在进行剪枝以后，效果并不明显。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
| class Solution {
private:
struct cmp{
bool operator()(pair<int,int>& x,pair<int,int>& y){
return x.first+x.second<y.first+y.second;
}
};
public:
vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
vector<pair<int,int>> result;
if(k<=0) return result;
priority_queue<pair<int,int>,vector<pair<int,int>>,cmp> pq;
for(int i=0;i<nums1.size();i++){
if(pq.size()>k){
auto t=pq.top();
if(nums1[i]+nums2[0]>=t.first+t.second) break;
}
for(int j=0;j<nums2.size();j++){
if(pq.size()>k){
auto t=pq.top();
if(nums1[i]+nums2[j]>=t.first+t.second) break;
}
pq.push(make_pair(nums1[i],nums2[j]));
if(pq.size()>k) pq.pop();
}
}
result.resize(pq.size());
for(int i=pq.size()-1;i>=0;i--){
result[i]=pq.top();
pq.pop();
}
return result;
}
};
|
根据博客 [LeetCode] Find K Pairs with Smallest Sums 找和最小的 K 对数字 中的解法四，可以得到如下答案，比之前要快不少。
具体原理是使用一个数组跟踪 nums1 中的每一位当前是和 nums2 中哪一位相加的，然后找出相加和最小的，下次 nums1 中的这一位就要和 nums2 中的下一位相加了。这样，本来 O(n^2)的复杂度就变成了 O(n*k)了（如果在最开始找到 nums1 和 nums2 中比较短的，能进一步优化）。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
| class Solution {
public:
vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
vector<pair<int,int>> result;
if(k==0 || nums1.size()==0 || nums2.size()==0) return result;
int most=min(int(nums1.size()*nums2.size()),k);
vector<int> trace(nums1.size(),0);
for(int i=0;i<most;i++){
int pos=0,sum=INT_MAX;
for(int j=0;j<nums1.size();j++){
if(trace[j]<nums2.size() && nums1[j]+nums2[trace[j]]<sum){
sum=nums1[j]+nums2[trace[j]];
pos=j;
}
}
result.push_back(pair<int,int>(nums1[pos],nums2[trace[pos]]));
trace[pos]++;
}
return result;
}
};
|
在 Solution 中找到了一种最快的方法Clean 16ms C++ O(N) Space O(KlogN) Time Solution using Priority queue，这个方法的复杂度只有 O(mlogn)。
方法中使用了 multimap，它的键是两个值的和，而值是一个 pair，表示两个值分别在 nums1 和 nums2 中的位置。
首先，将 nums1 中的每个元素与 nums2 中的第一个元素的和插入到 map 中，然后在循环中，每次取 map 中的第一个元素（也就是最小值），将其 push 到结果的 vector 中，然后将 nums1 中的当前元素和 nums2 中的后一个元素插入到 map 中（nums 中后一个元素比前一个大，所以只有前一个元素被输入到 vector 中之后，后面的元素才有机会被再次选中），这样循环 m 次，即可得到相应的结果。需要注意的是边界条件的判断，如果 nums2 中的元素 index 大于它的长度-1，就说明这个组合已经被用完了，再添加新的元素可能会越界。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
| class Solution {
private:
public:
vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
vector<pair<int,int>> result;
if(k<=0 || nums1.size()==0 || nums2.size()==0) return result;
multimap<int,pair<int,int>> m;
for(int i=0;i<nums1.size();i++) m.insert(make_pair(nums1[i]+nums2[0],make_pair(i,0)));
while(k-->0 && !m.empty()){
auto t=m.begin();
if(t==m.end()) break;
int a=t->second.first,b=t->second.second;
result.push_back(make_pair(nums1[a],nums2[b]));
m.erase(t);
if(b==nums2.size()-1) continue;
m.insert(make_pair(nums1[a]+nums2[b+1],make_pair(a,b+1)));
}
return result;
}
};
|