两序列和的最小 k 个数
题目:给定两个数组 a_1, \dots, a_n 和 b_1, \dots, b_m,从数组 a 和数组 b 分别取一个数然后相加,可以得到 nm 个数,找出这些数中的最小 k 个数。
数据范围:n, m, k \le 10^6。
思路
将 b 从小到大排序。然后将下标对 (i, 1) 插入到优先队列中,规定优先队列以 a_i + b_j 为键比较大小。
每次从优先队列取出当前的最小数 (i, j),然后再将 (i, j + 1) 插入到优先队列之中。
b 的排序保证了较小的数出现在较大的数前面。
时间复杂度为 O(m \log m + n + k \log n),其中 O(m \log m) 为排序,O(n) 为建堆。
代码
#include <iostream>
#include <algorithm>
#include <queue>
#include <utility>
using namespace std;
constexpr size_t MAXN = 1e5 + 5;
int a[MAXN], b[MAXN];
struct Pair {
size_t i, j;
friend bool operator<(Pair u, Pair v)
{
return a[u.i] + b[u.j] > a[v.i] + b[v.j] || \
(a[u.i] + b[u.j] == a[v.i] + b[v.j] && a[u.i] > a[v.i]);
}
};
int main()
{
ios::sync_with_stdio(false);
cin.tie(nullptr);
size_t n, m, k;
cin >> n >> m >> k;
for (size_t i = 1; i <= n; i++)
{
cin >> a[i];
}
for (size_t i = 1; i <= m; i++)
{
cin >> b[i];
}
sort(b + 1, b + m + 1);
priority_queue<Pair> qu;
for (size_t i = 1; i <= n; i++)
{
qu.push({i, 1});
}
for (size_t _ = 1; _ <= k; _++)
{
auto p = qu.top();
qu.pop();
size_t i = p.i, j = p.j;
cout << a[i] << ' ' << b[j] << '\n';
if (j < m)
{
qu.push({i, j + 1});
}
}
return 0;
}
BEYOND THE ARTICLE
读完了,有什么想法?
阅读已保存的问答示例,并循着引用回到原文。
— 本文完 — 返回文章列表