[洛谷] P1631 合并序列

题目描述

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

输入格式

第一行一个正整数N；

第二行N个整数$A_i$, 满足$A_i\le A_{i+1}$且$A_i\le 10^9$;

第三行N个整数$B_i$, 满足$B_i\le B_{i+1}$且$B_i\le 10^9$.

【数据规模】

对于50%的数据中，满足1<=N<=1000；

对于100%的数据中，满足1<=N<=100000。

输出格式

输出仅一行，包含N个整数，从小到大输出这N个最小的和，相邻数字之间用空格隔开。

输入输出样例

输入 #1

3
2 6 6
1 4 8

输出 #1

3 6 7

题目链接

https://www.luogu.com.cn/problem/P1631

题解

基础理解

由于题目满足$A_i\le A_{i+1}，B_i\le B_{i+1}$可以得到N个有序表：

$A[1]+B[1]<=A[1]+B[2]<=\dots <=A[1]+B[N]$

$A[2]+B[1]<=A[2]+B[2]<=\dots <=A[2]+B[N]$

$\dots \dots$

$A[N]+B[1]<=A[N]+B[2]<=\dots <=A[N]+B[N]$

即通过固定某个数组元素，对另一个数组元素进行遍历得到。

数据结构

优先队列：

priority_queue<pair<int, int>, vector<pair<int, int> >, greater<pair<int, int> > > q;

pair：

pair 绑定两个元素，第一个元素为 $first$，第二个元素为$second$。

在本题中，$first$表示两个数的和，如$A[1]+B[1]$，$second$ 表示当前元素在第几个有序表中。

思路

假设我们现在已经读取完两个数组。

上图中每一列表示一个有序表，我们将每个有序表的最小元素加入有序表：

for (int i = 1; i <= N; i++) {
    
    
        q.push(pair<int, int>(A[i] + B[1], i));
}

接下来开始输出，首先肯定会弹出$A[1]+B[1]$，然后接下来有可能会弹出哪个元素呢？没错，就是$A[2]+B[1]$或者$A[1]+B[2]$，我们现在已经把$A[2]+B[1]$加入了优先队列，所以现在我们要把$A[1]+B[2]$加入优先队列，让程序帮我们比较他俩的大小。

首次弹出的元素为$pair<A[1]+B[1],1>$, $1$就代表第一个有序表，即第一列，我们现在需要取出第一列的第二个元素入队列，但是我们肯定不能用$B[2]$来表示，万一下次需要取出$B[3]$呢，所以我们要记录每一个有序表加入了几个元素了。这里我们采用一个数组$from[i]$表示第 $i$ 个有序表加入了多少个元素。

修改前面的代码为：

for (int i = 1; i <= N; i++) {
    
    
        from[i] = 1;
        q.push(pair<int, int>(A[i] + B[1], i));
}

为了方便查看再放一次图：

我们在队列中加入了$A[1]+B[2]$， 接下来输出的是$A[2]+B[1]$或者$A[1]+B[2]$之一，假设是$A[2]+B[1]$ ，也就是说$A[2]+B[1]<A[1]+B[2]$，那么接着，我们会加入哪个元素入队列呢？比$A[2]+B[1]$ 大的$A[1]+B[2]$和$A[3]+B[1]$已经加入队列，还有可能比$A[2]+B[1]$大$(这里的大是指相邻的大，不是绝对大)$的元素是$A[2]+B[2]$，也就是说，每次从队列中取出一个元素，得到这个元素来自第几个有序表，再把这个有序表中未加入过优先队列的最小的元素加入优先队列实质是把有可能比当前元素值相邻大的元素入队列。

重复上述操作，直到输出$N$个元素为止。

完整代码如下：

#include <bits/stdc++.h>

using namespace std;

int main() {
    
    
    priority_queue<pair<int, int>, vector<pair<int, int> >, greater<pair<int, int> > > q;
    int N;
    cin >> N;
    int A[N+1], B[N+1];
    int from[N+1];
    //为方便理解，让下表从1开始
    for (int i = 1; i <= N; i++) {
    
    
        cin >> A[i];
    }
    for (int i = 1; i <= N; i++) {
    
    
        cin >> B[i];
        from[i] = 1;	//从每个有序表中各加入 1 个元素
        q.push(pair<int, int>(A[i] + B[1], i));
    }
    while (N--) {
    
    
        cout << q.top().first << " ";	//取得队头元素
        int next = q.top().second;	//第next个有序表
        q.pop();	
        //把第next个有序表中的第from[next]+1个元素入队
        q.push(pair<int, int>(A[next] + B[++from[next]], next));
    }
    return 0;
}