【洛谷】 P1631 序列合并

题目描述

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

输入格式

第一行一个正整数 $N$ ；
第二行 $N$ 个整数 $A_1,A_2,…,A_N$ 。
第三行 $N$ 个整数 $B_1,B_2,…,B_N$ 。

输出格式

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

样例输入 #1

3
2 6 6
1 4 8

样例输出 #1
3 6 7
提示

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

题解

对于输入的两个数组 $A=\{a_i | i=1,2,…N\},B=\{b_i | i=1,2,…N\}$ ，最直接的方法是用一个规格为 $N \times N$ 的数组去存放所有的 $a_i +b_j$ ，然后可采用基于选择排序的思路去选出前 $N$ 个最小的值。但是这样求解的时间复杂度为 $O(N^3 )$ ，这将难以通过全部测试数据。因此需要想别的办法。

如果你足够敏感，你应该会立刻想到一个数据结构——堆。堆是一种建立为 $O (n)$ 级、插入和删除都为 $O (l o g n)$ 级的数据结构。显然，在面对 $N \times N$ 个数据时，用他们直接建立规格为 $N \times N$ 的堆依然有可能超时（此时的时间复杂度为 $O(N^2)$ ），因此需考虑别的方法。试想，我们能否能利用堆的这种具有快速修改能力的数据结构来动态构建一个长度为 $N$ 的数据结构，并通过某种较为快速的取值方式对 $N \times N$ 个数据进行快速扫描并取值以更新该数据结构。为此，引入偏序集：

构建偏序集要求输入的数组有序。观察题中的两个数组： $a_1,a_2,…,a_N]，[b_1,b_2,…,b_N]$ ，它们的长度都为 $N$ ，对这两个数组分别排序的时间复杂度（最快）为 $O (Nl o g N)$ ，这是能接受的。当输入的两个数组均有序时，由所有的 $a_i +b_j$ 便可以组成以下 $N$ 个偏序集：
${ a_1+b_1, a_1+b_2, … ,a_1+b_N \} \\ \{ a_2+b_1, a_2+b_2, … ,a_2+b_N \} \\ …… \\ \{ a_N+b_1, a_N+b_2, … ,a_N+b_N \}$

显然，在这 $N$ 个偏序集中，其都是有序的（单调递增），因此，对每个单独的偏序集而言，其始终满足：
$a_i+b_j≤ a_i+b_{(j+1)}$

同时，可断言 $a_1+b_1$ 为所有 $a_i+b_j$ 组合中的最小值； $a_N+b_N$ 为所有 $a_i+b_j$ 组合中的最大值；每个偏序集中的最小值为 $a_i+b_1 (i=1,2,…,N)$ 、最大值为 $a_i+b_N (i=1,2,…,N)$ ；不同偏序集对应同一列的元素中始终有 $a_i+b_j≤ a_{(i+1)}+b_j$ 。

基于偏序集的这一特性，我们就能在 $O (1)$ 的时间复杂度内，从构建好的偏序集中取出当前最小值。具体的取值策略如下（假设现在有一个已经构建好的长度为 $N$ 的小根堆 $h e a p$ ）：

1、把每个偏序集中的最小元素加入小根堆 $h e a p$ 中，即 $heap=\{a_1+b_1,a_2+b_1,…,a_N+b_1\}$ ；
2、从小根堆 $h e a p$ 中取出根元素（即当前堆里的最小值），假设取出的元素为 $a_i+b_j$ ，记弹出数 +1；
3、从取出元素所在偏序集中，取出仅比其小的元素（即 $a_i+b_{(j+1)}$ ），将其插入小根堆 $h e a p$ 中；
4、若弹出数不为 $N$ ，则继续执行 2；否则结束取值。

当此算法结束时，我们就取出了所有 $a_i+b_j$ 组合中的前 $N$ 小值。

正确性证明

加入偏序集元素的顺序使得我们能保证当前加入的元素是该偏序集中最小的，而小根堆 $h e a p$ 又能保证每次取出的元素是 $N$ 个偏序集中未取的最小的元素中最小的，所以正确性得证。

STL 模板的使用

需要注意的是：在实际编码时，我们不必手动构建一个堆，而可以通过调用 STL 中的模板函数、数据结构来直接使用。优先队列priority_queue正是一个由 STL 提供的模板容器，该队列具有普通队列所具备的一切功能，不同的是，优先队列采用了“堆”这一数据结构，保证了该队列中的第一个元素总是整个优先队列中最大的（或最小的）元素（默认是大根堆）。
priority_queue 类的模板参数列表有三个参数：

class T，T 是优先队列中存储的元素的类型；
class Container = vector<T>，Container 是优先队列底层使用的存储结构，可以看出来，默认采用 vector；
class Compare = less<typename Container::value_type> Compar，是定义优先队列中元素比较方式的类。什么是优先队列中元素的比较方式？前面曾提到，优先队列的底层容器其实就是堆，而堆中元素是存在大小关系的。比如大根堆：每个结点的值都不大于它的父结点，因此其堆顶元素是最大的。当我们自己定义了某些类，并将由该类实例化得到的对象存进堆中时（此时堆里的元素就是某个对象），这时堆中元素的比较方式自然会发生改变。而编译器中的比较类只能比较内置类型，所以这时就需要用户给出一个用于比较大小的模式（即自定义类的比较方式），并将其填充到 Compare 参数的位置。

注：class Compare后跟着的less<typename Container::value_type>就是默认的比较类，默认是按小于（less）的方式比较，这种比较方式创建出来的是大根堆（所以优先队列默认为大根堆）。如果需要创建小堆，就需要将 less 改为 greater。

less<typename Container::value_type>的定义如下：

template <class T> 
struct less: binary_function <T,T,bool> {
    
    
	bool operator() (const T& x, const T& y) const {
    
    return x<y;}
};

greater<typename Container::value_type>的定义如下：

template <class T> 
struct greater: binary_function <T,T,bool> {
    
    
	bool operator() (const T& x, const T& y) const {
    
    return x>y;}
};

在这两个函数内部都重载了“( )”，从而构成仿函数以便使用。

方法一

前面的算法曾提到一点：当从优先队列中取出某个元素 $a_i+b_j$ 时，需要从该元素所在偏序集中取出仅次（小）于该元素的另一个元素 $a_i+b_{(j+1)}$ 填补进优先队列。理论上是这样进行取值，但实际操作时，我们不可能真的去建立一个 $N \times N$ 的二维数组（作为偏序集），因为这存在爆内存的可能。那要通过怎样的方式来完成这一功能呢？
回看对偏序集的取数操作，不难发现一点：当从优先队列中弹出某个元素 $a_i+b_j$ 时，该元素存在两个索引（即 $i$ 和 $j$ ），其中第一个索引 $i$ 指示了接下来要在 A 数组中取值的索引，第二个索引在执行 $j + 1$ 后指示了接下来要在 B 数组中取值的索引。因此，我们实际上只需要再建立 $N$ 个数对（为长度为N的优先队列中的每个元素各配置一个），就能完成从偏序集中取所需值的需求。
对于此结构（含：当前数的值 value、对应在A中索引 ida、对应在B中的索引 idb，共三个属性），因为他足够简单，因此我们可以自行构建结构体。但是需要注意一点：由于这个类型的数据会被放进优先队列中，因此需要为该结构指明一种比较大小的方式（当然，这里比较的肯定是 value），这就需要重载比较运算符；更进一步，我们要求此优先队列是“小根堆”，因此要重载的运算符是“>”。下面给出基于该思路写出的 AC 代码：

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

// 数据声明 
const int N = 100005;
int n,k,a[N],b[N];

// 自定义数据结构 
struct NODE {
    
    
    int value, ida, idb;
	
	// 重载运算符
    bool operator>(const NODE &obj) const {
    
     return value > obj.value;}
}node;

// 定义优先队列 
priority_queue<NODE,vector<NODE>,greater<NODE> > q;

int main(){
    
    
	
	// 录入数据 
    cin>>n;
    for (int i=0; i<n; i++) cin>>a[i];
    for (int i=0; i<n; i++) cin>>b[i];
    
    // 对数组排序 
    sort(a, a+n);
    sort(b, b+n);
    
	// 初始化优先队列 
    for (int i = 0; i<n; i++)
        q.push({
    
    a[i]+b[0], i, 0});
        
    while(n--){
    
    
    	// 取出当前队列的最小值 
        node=q.top(), q.pop();
        
        // 输出 
        cout<<node.value<<" ";
        
        // 从取出元素所在偏序集中选出其次元素（并更改索引 idb 值） 
        node.value = a[node.ida] + b[++node.idb];
        
        // 入队列 
        q.push(node);
    }
    return 0;
}

方法二

都说 STL 是个大暖男，这道题为此提供了充分的论据。
STL 提供了一个表达“数对”含义的数据结构 pair，其相关属性和函数定义如下：

构造函数pair<type1, type2>obj(value1, value2)：此函数可以声明并定义一个pair对象；
套娃函数make_pair( obj1, obj2)：此函数可以将两个已存在的对象构造为一个新的pair类型；
first属性：pair_obj.first 可取出对象 pair_ obj 中的第一个属性值；
second属性：pair_obj.second 可取出对象pair_ obj 中的第二个属性值。

根据前面的分析我们知道，现在需要一个含 3 个 int 类型组合在一起的数据结构，而 pair 结构体能够建立一个含两个元素的类。因此，为了能多加一个元素进去，我们可以再多建立一重（就像套娃一样再多一重）pair，即建立：pair<int, pair<int, int>> obj。这样一来，我们可设各层属性表达的含义如下：

obj.first // 表示 value
obj.second.first; // 表示 ida
obj.second.second // 表示 idb

如此，就能根据前面的思路写出求解本题的完整代码:

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

// 数据声明 
const int MAX = 100005;
int a[MAX], b[MAX];

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

int main(){
    
    
	int n, tmpa, tmpb;
	
	// 录入数据 
	cin>>n;
	for(int i=0;i<n;i++) cin>>a[i];
	for(int i=0;i<n;i++) cin>>b[i];
	
	// 对数组排序
	sort(a, a+n);
	sort(b, b+n); 
	
	for(int i=0;i<n;i++) q.push(make_pair(a[i]+b[0], pair<int,int>(i, 0)));
	
	// 查找前 k 小值 
	while(n--){
    
    
		// 输出 
		cout<<q.top().first<<" ";
		
		// 取出 ida 和 idb 索引 
		tmpa = q.top().second.first;
		tmpb = q.top().second.second;
		
		// 第一个元素出栈 
		q.pop();
		
		// 更改取出元素所在偏序集中的次元素入栈 
		q.push(make_pair(a[tmpa]+b[tmpb+1], pair<int,int>(tmpa, tmpb+1)));
	}
	return 0;
}