原题链接：POJ2442
给定 $m$ 组长度为 $n$ 的序列，每次从每一组中取出一个数，共有 $n^m$ 种取法，求所有取法中，取出的数字之和，前 $n$ 小是多少。
优先队列
考虑用两个优先队列来维护，一个小根堆 q，一个大根堆 heap。q 维护从前 i 组每组取一个数所能得到的前 n 小的和，heap 维护从前 i + 1 组取能得到的前 n 小的和。
初始化将第一组全部加入 q，然后枚举剩下的组，假设当前枚举的是第 i 组，那么就从 q 中弹出堆顶（最小值），与第 i 组的所有元素分别求和，加入大根堆 heap，同时保证 heap 中元素不超过 n 个，如果已经满了，如果当前取出的小根堆的堆顶与第 i 组某数求和后小于大根堆的堆顶，那么就可以舍弃大根堆的堆顶，再把该和值加入大根堆。处理完第 i 组后，将 q 重置为 heap 中的所有元素。
最后，小根堆 q 中存的就是前 n 小的和。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
|
#include <iostream>
#include <queue>
#include <cstdio>
#include <algorithm>
using namespace std;
typedef pair<int, int> PII;
typedef long long LL;
const int N = 2010, M = 110;
int T, n, m;
int a[M][N];
priority_queue<int> heap;
priority_queue<int, vector<int>, greater<int>> q;
int main() {
for (cin >> T; T--; ) {
cin >> m >> n;
for (int i = 1; i <= m; i++) {
for (int j = 1; j <= n; j++) {
scanf("%d", &a[i][j]);
}
}
for (int i = 1; i <= n; i++) q.push(a[1][i]);
for (int i = 2; i <= m; i++) {
while (q.size()) {
int t = q.top();
q.pop();
for (int j = 1; j <= n; j++) {
if (heap.size() < n) heap.push(t + a[i][j]);
else if (t + a[i][j] < heap.top()) {
heap.pop();
heap.push(t + a[i][j]);
}
}
}
while (heap.size()) {
q.push(heap.top());
heap.pop();
}
}
while (q.size()) {
printf("%d ", q.top());
q.pop();
}
puts("");
}
return 0;
}
|