[题解] Luogu 1631 – 序列合并

题目描述

有两个长度都是 $N$ 的序列 $A$ 和 $B$ ，在 $A$ 和 $B$ 中各取一个数相加可以得到 $N^2$ 个和，求这 $N^2$ 个和中最小的 $N$ 个。

输入格式

第一行一个正整数 $N$
第二行 $N$ 个整数 $A_i$，满足 $A_i \leq A_{i+1}$ 且 $A_i \leq 10^9$
第三行 $N$ 个整数 $B_i$，满足 $B_i \leq B_{i+1}$ 且 $B_i \leq 10^9$

输出格式

输出仅一行，包含 $N$ 个整数

样例输入

3
2 6 6
1 4 8

样例输出

3 6 7

解题思路

将数组 $A$ 和 数组 $B$ 排序，那么对于 $A$ 和 $B$ 相加后的结果可以分作 $n$ 组，并且满足

$$
\begin{aligned}
A_1 + B_1 \leq A_1 + B_2 & \leq \cdots \leq A_1 + B_n \\
A_2 + B_1 \leq A_2 + B_2 & \leq \cdots \leq A_2 + B_n \\
\cdots \\
A_n + B_1 \leq A_n + B_2 & \leq \cdots \leq A_n + B_n \\
\end{aligned}
$$

将每一组结果的第一个存入一个小根堆 $S$ 中，每一次在堆 $S$ 中取出最小的数 $x$ 并输出，同时将这个结果 $x$ 所在组的后一个数加入小根堆中即可

代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <queue>
#define P pair<int, int>

using namespace std;
const int MAXN = 100000 + 10;

inline int read(){
    int x = 0;
    int p = 1;
    char ch = getchar();

    while(ch < '0' || ch > '9'){
        if(ch == '-')
            p = 0;
        ch = getchar();
    }

    while('0' <= ch && ch <= '9'){
        x = x*10 + ch - '0';
        ch = getchar();
    }

    return p ? x : (-x);
}

int n;
int A[MAXN], B[MAXN], pos[MAXN];

priority_queue<P, vector<P>, greater<P> >que;
int main(){
    n = read();

    for(register int i=1; i<=n; i++){
        A[i] = read();
    }   

    for(register int i=1; i<=n; i++){
        B[i] = read();
    }

    sort(A+1, A+n+1);
    sort(B+1, B+n+1);

    for(register int i=1; i<=n; i++){
        pos[i] = 1;
        que.push(make_pair(A[i] + B[1], i));
    }

    for(register int i=1; i<=n; i++){
        P front = que.top();
        printf("%d ", front.first);
        que.pop();

        que.push(make_pair(A[front.second]+B[++pos[front.second]], front.second));
    }

    return 0;
}