【ACWing】146. 序列

题目地址：

https://www.acwing.com/problem/content/description/148/

给定$m$个序列，每个包含$n$个非负整数。现在我们可以从每个序列中选择一个数字以形成具有$m$个整数的序列。很明显，我们一共可以得到$n^m$个这种序列，然后我们可以计算每个序列中的数字之和，并得到$n^m$个值。现在请你求出这些序列和之中最小的$n$个值。

输入格式：
第一行输入一个整数$T$，代表输入中包含测试用例的数量。
接下来输入$T$组测试用例。
对于每组测试用例，第一行输入两个整数$m$和$n$。
接下在$m$行输入$m$个整数序列，数列中的整数均不超过$10000$。

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

数据范围：
$0<m\le 1000$
$0<n\le 2000$

如果$m=2$，是容易的，可以用最小堆来做，堆里存三元组$(y_1,y_2,s)$，$y_1,y_2$是两个序列分别取到的列下标，$s$是对应的两个数的和。堆顶是$s$最小的元组。先对两个序列从小到大排序，接着让$(1,1,s_0)$入堆，很显然这个$s_0$（即两个序列最小值之和）肯定是最小的。接着分别考虑$(2,1,.),(1,2,.)$，以此类推。注意相同元素不要重复考虑。

如果$m>2$，则可以两个两个考虑，每次只求前$n$大的和即可。代码如下：

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

const int N = 2010, M = 1010;
int m, n;
int a[M][N], c[N];
bool vis[N][N];
struct Node {
  int y1, y2, sum;
};
auto comp = [](Node& p1, Node& p2) { return p1.sum > p2.sum; };

int main() {
  int T;
  scanf("%d", &T);
  while (T--) {
    scanf("%d%d", &m, &n);
    int sum = 0;
    for (int i = 1; i <= m; i++) {
      for (int j = 1; j <= n; j++) scanf("%d", &a[i][j]);
      sort(a[i] + 1, a[i] + 1 + n);
    }
    for (int i = 2; i <= m; i++) {
      priority_queue<Node, vector<Node>, decltype(comp)> heap(comp);
      memset(vis, 0, sizeof(vis));
      heap.push({1, 1, a[1][1] + a[i][1]});
      for (int j = 1; j <= n; j++) {
        auto t = heap.top(); heap.pop();
        c[j] = t.sum;
        int y1 = t.y1, y2 = t.y2;
        // vis用来防止重复考虑
        if (y1 + 1 <= n && !vis[y1 + 1][y2]) {
          heap.push({y1 + 1, y2, a[1][y1 + 1] + a[i][y2]});
          vis[y1 + 1][y2] = true;
        }
        if (y2 + 1 <= n && !vis[y1][y2 + 1]) {
          heap.push({y1, y2 + 1, a[1][y1] + a[i][y2 + 1]});
          vis[y1][y2 + 1] = true;
        }
      }

      for (int j = 1; j <= n; j++) a[1][j] = c[j];
    }

    for (int i = 1; i <= n; i++) printf("%d ", a[1][i]);
    puts("");
  }
}

时间复杂度$O(mn\log n)$，空间$O(mn)$。