【UVA 11997 K Smallest Sums】优先级队列

来自《训练指南》优先级队列的例题。

题目链接：http://acm.hust.edu.cn/vjudge/problem/viewProblem.action?id=18702

题意：给定k个整数数组，各包含k个元素。在每个数组中取一个元素加起来，可以得到k^k个和，求这些和中最小的k个值（不去重）。

数据范围：k [2, 750]

思路：暴力枚举k^k不可取。

“先来看问题的简化版：给出两个长度为k的数组A和B，分别在A和B中任取一个数并相加，可以得到k^2个和，求这些和中最小的k个。”

首先把A, B从小到大排序。由于只求最小的前k个值，我们可以先假定B[0]与A的k个元素分别相加得到的元素集合S就是最小的k个和，然后每次取出S的最小值S_min保存到结果数组R中，并把S_min减去B[i]、加上B[i+1]后得到的值再放回集合S中，保持S的大小总为k。这样经过k轮取出-放回操作后，R中的k个元素就是结果。我们把这一操作过程叫作A,B的合并。合并一次的复杂度为O(klogk)

S需要k次插入和删除，每次只取最小值，因此数据结构选用完全二叉堆实现的优先级队列很合适。

现在由两个数组A,B扩展到k个数组A1~Ak，只需把A1看作A, 依次处理 i: 2~k，每次把Ai合并到A1中，这样经过k次两两合并后，A1中的k个数就是结果。总的复杂度O(k²logk)

 1 #include <cstdio>
 2 #include <algorithm>
 3 #include <queue>
 4 using namespace std;
 5 const int MAX_K = 760;
 6
 7 struct Elem{
 8     int sum, b;
 9     Elem(int s, int b):sum(s), b(b){}
10     Elem& operator = (const Elem& e){
11         sum = e.sum;
12         b = e.b;
13         return *this;
14     }
15     bool operator < (const Elem& e) const{
16         return sum > e.sum; //小顶堆
17     }
18 };
19
20 void merge(int* A, int* B, int n){
21     priority_queue<Elem> pq;
22     for(int i=0; i<n; i++){
23         pq.push(Elem(A[i] + B[0], 0));
24     }
25     for(int i=0; i<n; i++){
26         Elem e = pq.top();
27         A[i] = e.sum;
28         pq.pop();
29         e.sum = e.sum - B[e.b] + B[e.b+1];
30         e.b = e.b + 1;
31         pq.push(e);
32     }
33 //    for(int i=0; i<n; i++){
34 //        printf("%d ", A[i]);
35 //    }
36 //    printf("\n");
37 }
38
39 int k;
40 int a[MAX_K], b[MAX_K];
41
42 int main()
43 {
44     freopen("11997.txt", "r", stdin);
45     while(~scanf("%d", &k)){
46         for(int i=0; i<k; i++){
47             scanf("%d", &a[i]);
48         }
49         sort(a, a+k);
50         for(int i=1; i<k; i++){
51             for(int j=0; j<k; j++){
52                 scanf("%d", &b[j]);
53             }
54             sort(b, b+k);
55             merge(a, b, k);
56         }
57         printf("%d", a[0]);
58         for(int i=1; i<k; i++){
59             printf(" %d", a[i]);
60         }
61         printf("\n");
62     }
63     return 0;
64 }