首页 > 代码库 > UVA 11997 K Smallest Sums 优先队列 多路合并

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

  vjudge 上题目链接:UVA 11997

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

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

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

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

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

  表1:A1 + B1 <= A1 + B2 <= A1 + B3 <= ....
  表2:A2 + B<= A2 + B<= A2 + B3 <= ....

  ......

  表n:An + B<= An + B<= An + B3 <= ....

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

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

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