首页 > 代码库 > UVA 11997 K Smallest Sums 优先队列 多路合并
UVA 11997 K Smallest Sums 优先队列 多路合并
vjudge 上题目链接:UVA 11997
题意很简单,就是从 k 个数组(每个数组均包含 k 个正整数)中各取出一个整数相加(所以可以得到 kk 个结果),输出前 k 小的和。
这时训练指南上的一道题,这道题的简化版其实在 15 年的广东省省赛出现过,当时是以送分题的形式出现的,可我还是没能做出来,归根到底还是看书不够,接触的题型不够多。
*************************************************************大白书上的讲解开始***********************************************************************
在解决这个问题之前,先看看它的简化版:给出两个长度为 n 的有序表 A 和 B,分别在 A 和 B 中任取一个数并相加,可以得到 n2 个和。求这些和中最小的 n 个和。
这个问题可以转化为多路归并排序问题:即把 k 个有序表合并成一个有序表(假定每个表已经是升序排列)—— 用优先队列维护每个表的“当前元素”。如果一共有 n 个元素,则时间复杂度为 O(nlogk)。此时我们需要把这 n2 个和组织成如下 n 个有序表:
表1:A1 + B1 <= A1 + B2 <= A1 + B3 <= ....
表2:A2 + B1 <= A2 + B2 <= A2 + B3 <= ....
......
表n:An + B1 <= An + B2 <= An + B3 <= ....
其中第 a 张表里的元素形如 Aa + Bb 。用二元组 (s, b) 来表示一个元素,其中 s = Aa + Bb 。为什么不保存 A 的下标 a 呢?因为我们用不到 a 的值。如果我们需要得到一个元素 (s, b) 在表 a 中的下一个元素 (s‘, b+1),只需要计算 s‘ = Aa + Bb+1 = Aa + Bb - Bb + Bb+1 = s - Bb + Bb+1,并不需要知道 a 是多少。代码里可以用到如下结构体来表示。
struct Item { int s, b; // s = A[a] + B[b]。这里的 a 并不重要,因此不保存 Item(int s, int b): s(s), b(b) { } bool operator < (const Item &rhs) const { return s > rhs.s; } };
因为在任意时刻,优先队列中恰好有 n 个元素,一共取了 n 次最小值,因此时间复杂度为 O(nlogn)。代码如下:
//假设 A 和 B 的元素已经从小到大排序好 void merge(int *A, int *B, int *C, int n) { priority_queue<Item> q; for(int i = 0; i < n; ++i) q.push(Item(A[i] + B[0], 0)); for(int i = 0; i < n; ++i) { Item item = q.top(); q.pop(); // 取出 A[a] + B[b] C[i] = item.s; int b = item.b; if(b + 1 < n) q.push(Item(item.s - B[b] + B[b + 1], b + 1)); // 加入 A[a] + B[b + 1] = s - B[b] + B[b + 1] } }
而这题不是两个表,而是 k 个表,怎么办呢?两两合并就可以了(想一想,为什么),代码如下:
const int maxn = 768; int A[maxn][maxn]; int main() { int n; while(scanf("%d", &n) == 1) { for(int i = 0; i < n; ++i) { for(int j = 0; j < n; ++j) scanf("%d", &A[i][j]); sort(A[i], A[i] + n); } for(int i = 0; i < n; ++i) // 两两合并 merge(A[0], A[i], A[0], n); // (*) printf("%d",A[0][0]); // 输出结果 for(int i = 0; i < n; ++i) printf(" %d", A[0][i]); printf("\n"); } return 0; }
注意(*)处,merge 函数对 A[0] 又读又写,会有问题吗?(其实仔细观察函数就可以发现不会有任何影响,因为原来的值读完一次后就再也没用了,所以可以重复利用空间,便有了之后的写)。程序的复杂度为 O(k2logk)。另外,没有必要在一开始就把所有 k2 个元素保存在二维数组 A 中,而是可以每次只读 k 个元素,然后合并,从而大大降低空间复杂度。
*************************************************************大白书上的讲解结束***********************************************************************
以上是大白书上的讲解,根据它的思路我自己实现了本题的代码:
#include <cstdio> #include <cstring> #include <algorithm> #include <queue> using namespace std; #define For(i,s,t) for(int i = (s); i < (t); ++i) const int K = 760; int b[K],c[K]; struct Item { int sum, idx; Item(int _sum, int _idx): sum(_sum), idx(_idx) {} bool operator < (const Item &rhs) const { return sum > rhs.sum; } }; void _merge(int *b, int *c, int n) { priority_queue<Item> pq; For(i, 0, n) { pq.push(Item(b[i] + c[0], 0)); } For(i, 0, n) { Item Min = pq.top(); pq.pop(); b[i] = Min.sum; if(Min.idx + 1 < n) { pq.push(Item(Min.sum - c[Min.idx] + c[Min.idx + 1], Min.idx + 1)); } } } int main() { int k; while(~scanf("%d",&k)) { For(i, 0, k) { scanf("%d", b + i); } sort(b, b + k); For(p, 1, k) { For(i, 0, k) { scanf("%d", c + i); } sort(c, c + k); _merge(b, c, k); } printf("%d", b[0]); For(i, 1, k) { printf(" %d", b[i]); } puts(""); } return 0; }
感觉这道题蕴含的算法思想挺经典的,以后可能还会碰到,需要好好体会才行。
UVA 11997 K Smallest Sums 优先队列 多路合并