题目链接

算法竞赛进阶指南-60.序列

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

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

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

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

输入格式

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

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

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

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

输出格式

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

每组输出占一行。

数据范围

$0 < m \le 1000$,
$0 < n \le 2000$

输入样例：

1
2 3
1 2 3
2 2 3

输出样例：

3 3 4

Method : 优先队列

首先只考虑两个序列：$A = [a_1, a_2, a_3, ... a_n]$和 $B = [b_1, b_2, b_3, ... b_n]$
把数组A进行排序，得到$A' = [a'_1, a'_2, a'_3, ... a'_n]$，

分组：那么可以得到$n^2$个和，写成如下$n$组长度为$n$的序列，那么每个序列内部都是有序的（从小到大）
$[b_1 + a'_1, b_1 + a'_2, b_1 + a'_3, ... b_1 + a'_n]$

$[b_2 + a'_1, b_2 + a'_2, b_2 + a'_3, ... b_2 + a'_n]$

$[b_3 + a'_1, b_3 + a'_2, b_3 + a'_3, ... b_3 + a'_n]$

$[b_n + a'_1, b_n + a'_2, b_n + a'_3, ... b_n + a'_n]$
那么首先从这n组序列中选出最小的数，就是比较第一列的数，不妨假设最小的数是$b_2 + a'_1$，那么第二组的指针就向后移动，把$b_2 + a'_2$加进来继续比较，选出第二小的数…依次类推，可以得到两个序列的前n小和。
可以用一个堆来维护，时间复杂度为$O(n \log n)$

对于m个序列，我们仍然可以先只考虑两个序列，筛选出这两个序列的前n小和后，将其作为一个新的序列，继续与下一个序列求前n小和，一共需要求m - 1次，因此总时间复杂度为$O(mn \log n)$。

#include <queue>
#include <vector>
#include <iostream>
#include <algorithm>
#include <cstring>

using namespace std;

typedef pair<int, int> PII; // {当前的数, 当前数所在序列的下标}

const int N = 2000 + 7;

int T;
int m, n;
int a[N], b[N], tmp[N];

void merge() {
    priority_queue<PII, vector<PII>, greater<PII>> pq;
    for (int i = 0; i < n; i ++) {
        pq.push({a[0] + b[i], 0});
    }
    
    for (int i = 0; i < n; i ++) {
        auto t = pq.top(); pq.pop();
        tmp[i] = t.first;
        pq.push({t.first + a[t.second + 1] - a[t.second], t.second + 1});
    }
    
    // copy
    memcpy(a, tmp, sizeof tmp);
}

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

复杂度分析

• 时间复杂度：${O(mn\log n)}$， 其中$n\log n$是$n$次堆的push() & pop()，需要进行$m - 1$次求前n小和操作。

• 空间复杂度：${O(1)}$。