P1631 序列合并

P1631 序列合并

题目描述

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

输入格式

第一行一个正整数 $N$；

第二行 $N$ 个整数 $A_{1\dots N}$。

第三行 $N$ 个整数 $B_{1\dots N}$。

输出格式

一行 $N$ 个整数，从小到大表示这 $N$ 个最小的和。

输入输出样例 #1

输入 #1

3
2 6 6
1 4 8

输出 #1

3 6 7

说明/提示

对于 $50\%$ 的数据，$N\le 10^3$。

对于 $100\%$ 的数据，$1\le N\le 10^5$，$1\le a_i,b_i\le 10^9$。

题解

本题可以将所有序列的和都算出来，然后排序，输出前 $N$ 个。

但是这样的时间复杂度为 $O(N^2)$，会超时。

所以我们需要优化。

本题的关键在于单调不降。

那么显然，对于合并即相加后的序列，其一定是单调不降的。

即：

• a[1] + b[1] ≤ a[1] + b[2] ≤ a[1] + b[3] ≤ … ≤ a[1] + b[n] ；
• a[2] + b[1] ≤ a[2] + b[2] ≤ a[2] + b[3] ≤ … ≤ a[2] + b[n] ；

我们可以列表（横向b递增，竖向a递增）：

a[1] + b[1]	a[1] + b[2]	a[1] + b[3]	············	a[1] + b[n]
a[2] + b[1]	a[2] + b[2]	a[2] + b[3]	············	a[2] + b[n]
a[3] + b[1]	a[3] + b[2]	a[3] + b[3]	············	a[3] + b[n]
············	············	············	············	············
a[n] + b[1]	a[n] + b[2]	a[n] + b[3]	············	a[n] + b[n]

对于整个表，最小的一项一定是a[1] + b[1]，次小的一定是a[1] + b[2]或a[2] + b[1]，以此类推。

那我们可以用小根堆来维护这个表。

每次取出最小的一项，然后将其下一项加入队列。

这样的时间复杂度为 $O(NlogN)$。

具体来说，根据题意我们需要找出前 $N$ 小的和，所以我们需要维护一个大小为 $N$ 的优先队列。

那么我们可以先计算出第一行，即a[1] + b[1]到a[1] + b[n]，然后将其加入优先队列。

然后我们可以取出队列中的最小值，即a[1] + b[1]，然后将其下一项a[2] + b[1]加入队列。

代码实现

#include <iostream>
#include <queue>
using namespace std;

const int N = 1e5 + 10;

int n;
int a[N],b[N];

struct node{
	int sum;
	int i,j;
	bool operator < (const node& x) const{
		return sum > x.sum;
	}
};

priority_queue<node> heap;

int main(){
	cin >> n;
	for(int i = 1;i <= n;i++) cin >> a[i];
	for(int i = 1;i <= n;i++) cin >> b[i];
	
	for(int i = 1;i <= n;i++){
		heap.push({a[i] + b[1],i,1});
	}
	
	for(int k = 1;k <= n;k++){
		auto t = heap.top();heap.pop();
		int sum = t.sum,i = t.i,j = t.j;
		cout << sum << " ";
		if(j <= n) heap.push({a[i] + b[j + 1],i,j + 1});
	}
	return 0;
}