两个有序数组间相加和的Topk问题
题目描述
给定两个有序数组arr1和arr2,再给定一个整数k,返回来自arr1和arr2的两个数相加和最大的前k个,两个数必须分别来自两个数组
按照降序输出
[要求]
时间复杂度为 O ( k log k ) O(k \log k) O(klogk)
输入描述:
第一行三个整数N, K分别表示数组arr1, arr2的大小,以及需要询问的数
接下来一行N个整数,表示arr1内的元素
再接下来一行N个整数,表示arr2内的元素
输出描述:
输出K个整数表示答案
示例1
输入
5 4
1 2 3 4 5
3 5 7 9 11
输出
16 15 14 14
备注:
1
⩽
N
⩽
1
0
5
1 \leqslant N \leqslant 10^5
1⩽N⩽105
0
⩽
a
r
r
1
i
,
a
r
r
2
i
⩽
1
0
9
0 \leqslant arr_{1_{i}}, arr_{2_{i}} \leqslant 10^9
0⩽arr1i,arr2i⩽109
保证
1
⩽
K
⩽
2
N
1 \leqslant K \leqslant 2N
1⩽K⩽2N
题解:
一般涉及到 TopK
问题,首先应该想到使用堆2333。
我们可以知道 a[n-1] + b[n-1] 是整体的最大值,那么在这个最大值之后应该是哪个元素呢,应该是: m a x { a [ n − 1 ] + b [ n − 2 ] , a [ n − 2 ] + b [ n − 1 ] } max\{a[n-1] + b[n-2], a[n-2] + b[n-1]\} max{a[n−1]+b[n−2],a[n−2]+b[n−1]},它俩就是下一个可能最大值,绝不会出现比它们更大的元素。
我们可以使用一个大根堆,每个节点保存三个元素 { s u m , i d x 1 , i d x 2 } \{sum, idx1, idx2\} {sum,idx1,idx2},且 a [ i d x 1 ] + b [ i d x 2 ] = s u m a[idx1]+b[idx2] = sum a[idx1]+b[idx2]=sum,两个下标用来表明该元素由两个数组中哪两个元素凑成的。
选择堆顶元素,其就是整体的最大值,将其删除,然后把两个可能的最大元素插入到堆中去:
-
{ a [ i d x 1 − 1 ] + b [ i d x 2 ] , i d x 1 − 1 , i d x 2 } \{a[idx1-1]+b[idx2], idx1-1, idx2\} {a[idx1−1]+b[idx2],idx1−1,idx2}
-
{ a [ i d x 1 ] + b [ i d x 2 − 1 ] , i d x 1 , i d x 2 − 1 } \{a[idx1] + b[idx2-1], idx1, idx2-1\} {a[idx1]+b[idx2−1],idx1,idx2−1}
这样进行 k 次操作后,就可以得到整体的 TopK 了。
注意:
1. 需要判重,因为两个节点的下标有可能相等
2. 测试数据出现无序的情况。。。。。。
代码:
#include <cstdio>
#include <queue>
#include <algorithm>
#include <unordered_map>
using namespace std;
const int N = 100000;
typedef pair<int, int> PII;
typedef pair<int, PII> PIP;
int n, k;
int a[N], b[N];
int main(void) {
scanf("%d%d", &n, &k);
for ( int i = 0; i < n; ++i )
scanf("%d", a + i);
for ( int i = 0; i < n; ++i )
scanf("%d", b + i);
sort( a, a + n );
sort( b, b + n );
priority_queue<PIP> heap;
unordered_map<int, bool> hash;
heap.push({a[n-1]+b[n-1], {n-1,n-1}});
hash[(n - 1) * N + n - 1] = true;
while ( k-- ) {
auto t = heap.top();
heap.pop();
printf("%d ", t.first);
int x = t.second.first, y = t.second.second;
int key = x * N + y - 1;
if ( y && !hash.count(key) ) {
heap.push({a[x] + b[y-1], {x, y-1}});
hash[key] = true;
}
key = (x - 1) * N + y;
if ( x && !hash.count(key) ) {
heap.push({a[x-1] + b[y], {x-1, y}});
hash[key] = true;
}
}
return 0;
}
贴一个正儿八经的 O(klogk) 复杂度代码:
#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
using namespace std;
const int N = 1e5 + 10;
int a[N];
int b[N];
using PII = pair<int, int>;
int main(void) {
int n, k;
scanf("%d%d", &n, &k);
for ( int i = 0; i < n; ++i )
scanf("%d", a + i);
for ( int i = 0; i < n; ++i )
scanf("%d", b + i);
sort(a, a + n);
sort(b, b + n);
auto cmp = [&](const PII& lhs, const PII& rhs) {
return a[lhs.first] + b[lhs.second] <= a[rhs.first] + b[rhs.second];
};
priority_queue<PII, vector<PII>, decltype(cmp)> heap(cmp);
int r = max(0, n - k);
for ( int i = n - 1; i >= r ; --i ) {
heap.emplace( move(PII(i, n - 1)) );
}
while ( k-- ) {
auto [x, y] = heap.top();
heap.pop();
printf("%d ", a[x] + b[y]);
if ( y ) heap.emplace( move(PII(x, y-1)) );
}
return 0;
}