计算机

两序列和的最小 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

读完了,有什么想法?

阅读已保存的问答示例,并循着引用回到原文。

查看公开问答示例
— 本文完 — 返回文章列表