AcWing 146. 序列(堆 + 多路归并)

题目描述

给定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

题解

1. 将m个序列两两合并·，即先求出两个序列的前n个最小值。再用这n个最小值与下一组m序列求出新的n个最小值，重复这个过程，知道m个序列比完。
2. 考虑如何从两个序列中找到n个最小值：将这两个序列分别排序，然后按照题目中的意思进行分组，可以找到n组， 每组n个结果，即a0 + b1, a1 + b1 ,…, a0+b2, a1+ b2, … a0 + bn-1, a1 + bn-1.
   这样，a, b 都是有序的，因此，n组的每组第一个数一定是前n个最小值。取出这n个数最小值（可以用堆来维护）。然后删除。并且把删除的数的后一个数加入到堆中。要找到后一个数需要知道这个数的下标。

代码

#include <iostream>
#include <queue>
#include <algorithm>
#include <vector>
#include <cstring>
using namespace std;

typedef pair<int, int> PII;
const int N = 2010, M = 2010;
int a[N], b[N], c[N];
int n, m;

void work()
{
    priority_queue<PII, vector<PII>, greater<PII>> heap;
    for(int i = 0; i < n; i++) heap.push({a[0] + b[i], 0}); //加入n组中每组的第一个数
    
    for(int i = 0; i < n; i++) //取出前n个最小值
    {
        auto t = heap.top();
        heap.pop();
        c[i] = t.first;
        heap.push({t.first - a[t.second] + a[t.second + 1], t.second + 1}); //t.first - a[t.second] = b[i]
    }
    
    memcpy(a, c, 4 * n);//更新a数组，用来与下一组序列比较
}

int main()
{
    int T;
    scanf("%d", &T);
    while(T--)
    {
        scanf("%d%d", &m, &n);
        for(int i = 0; i < n; i++) cin >> a[i];
        sort(a, a + n);
        
        for(int i = 1; i < m; i++)
        {
            for(int j = 0; j < n; j++) cin >> b[j];
            sort(b, b + n);
            work();
        }
        
        for(int i = 0; i < n; i++) cout << a[i] << " ";
        
        cout << endl;
    }
    
    return 0;
}