VOJ n个最小和 题解 小顶堆

n个最小和

题目描述

给出两个包含 n 个整数的数组A，B。分别在A，B 中任意出一个数并且相加，可以得到 $n^2$ 个和。求这些和中最小的 n 个。

输入描述

输入第一行一个整数 n 接下来一行输入数组 A，用空格隔开。接下来一行输入数组 B，用空格隔开。

输出描述

从小到大输出最小的 n 个和,用空格隔开。

用例输入 1

4
1 3 5 7
2 4 6 8

用例输出 1

3 5 5 7

数据规模与约定

$1<=n<=50000$，
$1<=A_i,B_i<=10^9$。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

struct node
{
    int a_idx, b_idx, sum;              // a中索引，b中索引，两数相加和
    bool operator>(const node &x) const // 结构体类型的优先队列需要重载运算符
    {
        return sum > x.sum;
    }
};

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    int n;
    cin >> n;
    vector<int> a(n), b(n);
    for (int i = 0; i < n; i++)
    {
        cin >> a[i];
    }
    for (int i = 0; i < n; i++)
    {
        cin >> b[i];
    }
    // 一定要先排序，才能保证a[i+1]+b[j]>=a[i]+b[j]
    sort(a.begin(), a.end());
    sort(b.begin(), b.end());
    priority_queue<node, vector<node>, greater<node>> pq; // 小顶堆
    // 先将a0+bi共n个数压入栈中，并存储相对应的在a和b数组中的索引
    for (int i = 0; i < n; i++)
    {
        pq.push({0, i, a[0] + b[i]});
    }
    int cnt = n; // 需输出的和的个数
    while (cnt--)
    {
        node x = pq.top(); // 每次取出最小的和并输出结果
        pq.pop();
        cout << x.sum << ' ';
        if (x.a_idx != n - 1) // 如果a数组中的索引不是最后一个，则将下一个索引对应的a数组中的值和bi相加后加入优先队列中
        {
            pq.push({x.a_idx + 1, x.b_idx, a[x.a_idx + 1] + b[x.b_idx]});
        }
    }

    return 0;
}