题意:有K个整数数组,各包含K个元素。在每个数组中取一个元素加起来,可以得到k^k个和。求这些和中最小的K个值。
tags:思维,大白书189
简化版:两个数组 A[]、 B[],有 k*k 个和,怎么快速求出前 k 小和 ?
数组排序后,对于 s1= A[i]+B[j] ,它的下一个最大的和就是 A[i]+B[j+1] = s1-B[j]+B[j+1] 。 所以优先队列优化一下,就可以实现 O(n*log(n)) 。
对于 k 个数组,我们只要两两合并就可以了。
#include<bits/stdc++.h> using namespace std; #pragma comment(linker, "/STACK:102400000,102400000") #define rep(i,a,b) for (int i=a; i<=b; ++i) #define per(i,b,a) for (int i=b; i>=a; --i) #define mes(a,b) memset(a,b,sizeof(a)) #define INF 0x3f3f3f3f #define MP make_pair #define PB push_back #define fi first #define se second typedef long long ll; const int N = 805; int k, a[N][N]; struct Item { int id, sum; bool friend operator<(Item ca, Item cb) { return ca.sum > cb.sum; } }; priority_queue< Item > q; void Merge(int* A, int* B) { while(!q.empty()) q.pop(); rep(i,1,k) q.push((Item){ 1, A[i]+B[1] }); rep(i,1,k) { Item u=q.top(); q.pop(); A[i]=u.sum; if(u.id<k) q.push((Item){ u.id+1, u.sum-B[u.id]+B[u.id+1] }); } } int main() { while(~scanf("%d", &k)) { rep(j,1,k) scanf("%d", &a[1][j]); sort(a[1]+1, a[1]+1+k); rep(i,2,k) { rep(j,1,k) scanf("%d", &a[i][j]); sort(a[i]+1, a[i]+1+k); Merge(a[1], a[i]); } rep(j,1,k) printf("%d%c", a[1][j], " \n"[j==k]); } return 0; }