UVA 11997 K Smallest Sums 优先队列 多路合并

　　vjudge 上题目链接：UVA 11997

　　题意很简单，就是从 k 个数组（每个数组均包含 k 个正整数）中各取出一个整数相加（所以可以得到 k^k 个结果），输出前 k 小的和。

　　这时训练指南上的一道题，这道题的简化版其实在 15 年的广东省省赛出现过，当时是以送分题的形式出现的，可我还是没能做出来，归根到底还是看书不够，接触的题型不够多。

*************************************************************大白书上的讲解开始***********************************************************************

　　在解决这个问题之前，先看看它的简化版：给出两个长度为 n 的有序表 A 和 B，分别在 A 和 B 中任取一个数并相加，可以得到 n² 个和。求这些和中最小的 n 个和。

　　这个问题可以转化为多路归并排序问题：即把 k 个有序表合并成一个有序表（假定每个表已经是升序排列）—— 用优先队列维护每个表的“当前元素”。如果一共有 n 个元素，则时间复杂度为 O(nlogk)。此时我们需要把这 n² 个和组织成如下 n 个有序表：

　　表1：A₁ + B₁ <= A₁ + B₂ <= A₁ + B₃ <= ....
　　表2：A₂ + B₁ <= A₂ + B₂ <= A₂ + B₃ <= ....

　　......

　　表n：A_n + B₁ <= A_n + B₂ <= A_n + B₃ <= ....

　　其中第 a 张表里的元素形如 A_a + B_b 。用二元组 (s, b) 来表示一个元素，其中 s = A_a + B_b 。为什么不保存 A 的下标 a 呢？因为我们用不到 a 的值。如果我们需要得到一个元素 (s, b) 在表 a 中的下一个元素 (s‘, b+1)，只需要计算 s‘ = A_a + B_b+1 = A_a + B_b - B_b + B_b+1 = s - B_b + B_b+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(k²logk)。另外，没有必要在一开始就把所有 k² 个元素保存在二维数组 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;
}

　　感觉这道题蕴含的算法思想挺经典的，以后可能还会碰到，需要好好体会才行。