K Smallest Sums(Uva 11997) 多路归并+优先队列

来自《算法竞赛入门经典训练指南》

1.题目原文

https://uva.onlinejudge.org/index.php?option=com_onlinejudge&Itemid=8&page=show_problem&problem=3148

有k个整数数组，各包含k个整数。在每个数组中去一个元素加起来，可以得到k^k个元素，求这些和中最小的k个值(重复的值算多次)

2.解题思路

考虑题目的简化版本：

有两个长度为n的有序数组A和B。分别在A和B中去一个元素，把它们加起来，可以得到n^n个元素，求最小的n个值。

这个问题就转化为多路归并问题。

表1:A[1]+B[1]<=A[1]+B[2]<=A[1]+B[3]<=……<=A[1]+B[n];

表2:A[2]+B[1]<=A[2]+B[2]<=A[2]+B[3]<=……<=A[2]+B[n];

……

表n:A[n]+B[1]<=A[n]+B[2]<=A[n]+B[3]<=……<=A[n]+B[n];

其中表a中的元素形如A[a]+B[b]。我们用一个二元组(s,b)来表示一个元素，其中s=A[a]+B[b]。就可以得到元素(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。

表示元素(s,b)的结构体如下

struct Item
{
    int s,b;//s=A[a]+B[b];
    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的元素已经从小到大排好序
//数组C中存放n^2个和中最小的n个和
void merge(int *A,int *B,int *C,int n)
{
    priority_queue<Item> que;
    for(int i=0;i<n;i++){
        que.push(Item(A[i]+B[0],0));
    }
    for(int i=0;i<n;i++){
        Item item=que.top();
        que.pop();
        C[i]=item.s;
        int b=item.b;
        if(b+1<n) que.push(Item(item.s-B[b]+B[b+1],b+1));
        //加入A[a]+B[b+1]=s-B[b]+B[b+1]
    }
}

然后回到本题，共有k个数组，两两合并即可。时间复杂度为O(k^2logk)。

3.AC代码

#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <string>
#include <set>
#include <vector>
#include<cmath>
#include<bitset>
#include<stack>
#include<sstream>
using namespace std;
#define INF 0x7fffffff
#define maxn 1005

int n;
int A[maxn][maxn];

struct Item
{
    int s,b;//s=A[a]+B[b];
    Item(int s,int b):s(s),b(b){}
    bool operator < (const Item& rhs) const{
        return s>rhs.s;
    }
};

//假设A和B的元素已经从小到大排好序
//数组C中存放n^2个和中最小的n个和
void merge(int *A,int *B,int *C,int n)
{
    priority_queue<Item> que;
    for(int i=0;i<n;i++){
        que.push(Item(A[i]+B[0],0));
    }
    for(int i=0;i<n;i++){
        Item item=que.top();
        que.pop();
        C[i]=item.s;
        int b=item.b;
        if(b+1<n) que.push(Item(item.s-B[b]+B[b+1],b+1));
        //加入A[a]+B[b+1]=s-B[b]+B[b+1]
    }
}

void solve()
{
    for(int i=0;i<n;i++){
        sort(A[i],A[i]+n);
    }
    for(int i=1;i<n;i++){
        merge(A[0],A[i],A[0],n);
    }
    for(int i=0;i<n-1;i++){
        printf("%d ",A[0][i]);
    }
    printf("%d\n",A[0][n-1]);
}

int main()
{
    while(scanf("%d",&n)!=EOF){
        for(int i=0;i<n;i++){
            for(int j=0;j<n;j++){
                scanf("%d",&A[i][j]);
            }
        }
        solve();
    }
    return 0;
}