算法提高——数据结构(多路归并+二叉堆)

---

一、例题

序列

给定 m 个序列，每个包含 n 个非负整数。

现在我们可以从每个序列中选择一个数字以形成具有 m 个整数的序列。

很明显，我们一共可以得到 nm 个这种序列，然后我们可以计算每个序列中的数字之和，并得到 nm 个值。

现在请你求出这些序列和之中最小的 n 个值。

输入格式
第一行输入一个整数 T，代表输入中包含测试用例的数量。

接下来输入 T 组测试用例。

对于每组测试用例，第一行输入两个整数 m 和 n。

接下在 m 行输入 m 个整数序列，数列中的整数均不超过 10000。

输出格式
对于每组测试用例，均以递增顺序输出最小的 n 个序列和，数值之间用空格隔开。

每组输出占一行。

数据范围
0<m≤1000,
0<n≤2000
输入样例：
1
2 3
1 2 3
2 2 3
输出样例：
3 3 4

题目来源：算法竞赛进阶指南
题目链接：Acwing

二、思路分析

分析题目，题目让我们从m个序列中，寻找出和最小的n个数，如果用暴力去解的话，时间复杂度为O(n^m)，显然不可取。

对于这种问题，我们可以用多路归并来解决。题目的突破口是只需要寻找到和最小的n个数，假设我们只有两个序列，那我们做法就很简单了，只需要将全部的和算出来，只保留把这两个序列合成的最小n个数的和即可，这合成的新的序列就是答案。
那我们根据这个思路，其实也可以推广到m个序列的情况，因为m个序列最终都是要合并的，那我们不妨可以先任取其中两个合并，只保留最小的n个数，那在我们求出来的最终答案里，肯定是由这最小n个数递推而来的。因此，我们就解决了序列合并的问题了，但是这样的解法似乎还远远不够。如果两个序列我们把全部的和都算出来，时间复杂度就为O(n^2)，在加上m个序列，还是会超时。那我们此时再来优化一下两个序列合并之间的细节。

假设序列a和序列b合并，并且两个序列都是单调递增，那么此时我们可以试着将两个序列之间的组合分组。
分成
b1 + a1 ,b1 + a2,b1 + a3 … b1 + an
b2 + a1 ,b2 + a2,b2 + a3 … b2 + an
…
bn + a1 ,bn + a2,bn + a3 … bn + an
我们观察上面的序列，因为a和b序列都是单调递增的，所以每一组最左边必然为该组的最小值，而每次我们只需要比较n个数大小，从中取出最小的值，在更新取出元素该组的最小值。说到这里大家想必都反应过来了，这个操作不就是堆的操作吗。所以，这里我们就可以用优先队列来维护最小的n个数，优先队列的时间复杂度为logn,那我们的操作就从O(n^2)的序列变成了O(nlogn)，时间没问题，那接下来就是ac代码了~

三、代码

#include <cstdio>
#include <queue>
#include <algorithm>

using namespace std;
typedef pair<int,int> PII;
const int N = 2010;

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

void merge()
{
    priority_queue<PII,vector<PII>,greater<PII>> q;
    
    for (int i = 0; i < n; i ++ ) q.push({a[0] + b[i],0});
    
    for (int i = 0; i < n; i ++ )
    {
        PII t = q.top();  q.pop();
        c[i] = t.first;
        q.push({a[t.second + 1] + t.first - a[t.second] ,t.second + 1});
    }
    
    for (int i = 0; i < n; i ++ ) a[i] = c[i];
    
}

int main()
{
    int test;
    scanf("%d",&test);
    while(test -- )
    {
        scanf("%d%d",&m,&n);
        for (int i = 0; i < n; i ++ ) scanf("%d",&a[i]);
        
        sort(a,a + n);
        for (int i = 0; i < m - 1; i ++ )
        {
            for (int j = 0; j < n; j ++ ) scanf("%d",&b[j]);
            merge();
        }
        
        for (int i = 0; i < n; i ++ ) printf("%d ",a[i]);
        printf("\n");
    }
    
    return 0;
}