来自《训练指南》优先级队列的例题。
题目链接:http://acm.hust.edu.cn/vjudge/problem/viewProblem.action?id=18702
题意:给定k个整数数组,各包含k个元素。在每个数组中取一个元素加起来,可以得到kk个和,求这些和中最小的k个值(不去重)。
数据范围:k [2, 750]
思路:暴力枚举k^k不可取。
“先来看问题的简化版:给出两个长度为k的数组A和B,分别在A和B中任取一个数并相加,可以得到k^2个和,求这些和中最小的k个。”
首先把A, B从小到大排序。由于只求最小的前k个值,我们可以先假定B[0]与A的k个元素分别相加得到的元素集合S就是最小的k个和,然后每次取出S的最小值Smin保存到结果数组R中,并把Smin减去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(k2logk)
#include <cstdio>
#include <algorithm>
#include <queue>
using namespace std;
const int MAX_K = ; struct Elem{
int sum, b;
Elem(int s, int b):sum(s), b(b){}
Elem& operator = (const Elem& e){
sum = e.sum;
b = e.b;
return *this;
}
bool operator < (const Elem& e) const{
return sum > e.sum; //小顶堆
}
}; void merge(int* A, int* B, int n){
priority_queue<Elem> pq;
for(int i=; i<n; i++){
pq.push(Elem(A[i] + B[], ));
}
for(int i=; i<n; i++){
Elem e = pq.top();
A[i] = e.sum;
pq.pop();
e.sum = e.sum - B[e.b] + B[e.b+];
e.b = e.b + ;
pq.push(e);
}
// for(int i=0; i<n; i++){
// printf("%d ", A[i]);
// }
// printf("\n");
} int k;
int a[MAX_K], b[MAX_K]; int main()
{
freopen("11997.txt", "r", stdin);
while(~scanf("%d", &k)){
for(int i=; i<k; i++){
scanf("%d", &a[i]);
}
sort(a, a+k);
for(int i=; i<k; i++){
for(int j=; j<k; j++){
scanf("%d", &b[j]);
}
sort(b, b+k);
merge(a, b, k);
}
printf("%d", a[]);
for(int i=; i<k; i++){
printf(" %d", a[i]);
}
printf("\n");
}
return ;
}