洛谷题单指南-二叉堆与树状数组-P1631 序列合并

原题链接：https://www.luogu.com.cn/problem/P1631

题意解读：两个长度为n的单调不降的有序序列，各取一个数相加，计算最小的n个和。

解题思路：

设两个序列为a[N], b[N]

和最小的为a[1] + b[1]

和第二小的可能有两种：a[1] + b[2]或者a[2]+b[1]

假设第二小为a[1] + b[2]，那么和第三小的将在a[2] + b[2]，a[1] + b[3]，a[2] + b[1]中选

假设第三小为a[2] + b[1]，那么和第四小的将在a[2] + b[2]，a[1] + b[3], a[2] + b[2], a[3] + b[1]中选

注意a[2] + b[2]出现两次，应该去重，只在a[2] + b[2]，a[1] + b[3],  a[3] + b[1]中选

因此，此题算法逻辑如下：

1、设ai=1，bi=1，sum=a[ai]+b[bi]，将{ai, bi, sum}存入优先队列，sum小的在堆顶

2、取堆顶元素，输出sum

3、基于堆顶元素的ai，bi，扩展出{ai, bi + 1, a[ai] + b[bi + 1]}，{ai + 1, bi, a[ai + 1], b[bi]}，加入优先队列

注意：加入优先队列前要对已经加入过的下标进行去重，可以借助map<int, map<int, bool>>来对两个下标进行去重。

4、重复步骤2，一共n次

100分代码：

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

const int N = 100005;
int n;
int a[N], b[N];
map<int, map<int, bool>> flag;

struct Node
{
    int ai, bi, sum;
    bool operator < (const Node &node) const &
    {
        return sum > node.sum;
    }
};

priority_queue<Node> q;

int main()
{
    cin >> n;
    for(int i = 1; i <= n; i++) cin >> a[i];
    for(int i = 1; i <= n; i++) cin >> b[i];
    q.push({1, 1, a[1] + b[1]});
    for(int i = 1; i <= n; i++)
    {
        Node node = q.top(); q.pop();
        cout << node.sum << " ";
        if(!flag[node.ai + 1][node.bi])
        {
            q.push({node.ai + 1, node.bi, a[node.ai + 1] + b[node.bi]});
            flag[node.ai + 1][node.bi] = true;
        }
            
        if(!flag[node.ai][node.bi + 1])
        {
            q.push({node.ai, node.bi + 1, a[node.ai] + b[node.bi + 1]});
            flag[node.ai][node.bi + 1] = true;
        }
    }
    return 0;
}