UVA11997 K Smallest Sums 優先隊列

題意

有 k 個整數數組，各包含 k 個元素，在每個數組中取一個元素加起來，可以得到  個和。求這些和中最小的 k 個值（允許重複）。

分析

這個問題可以化簡成兩個長度爲 k 的有序表各取一個元素求前 k 個最小的。

拿二元組（s, b）來表示一個元素，，這樣訪問下一個元素時，訪問的是

代碼

#include <cstdio>
#include <iostream>
#include <algorithm>
#include <queue>
using namespace std;
const int maxn = 750 + 7;
int a[maxn];
int b[maxn];

struct Item{
    int s, b;
    Item(int s, int b): s(s), b(b) {}
    bool operator < (const Item& a) const{
        return a.s < s;
    }
};

void merge(int* a, int *b, int *c, int n){
    priority_queue<Item> pq;
    for(int i = 0; i < n; i++){
        pq.push(Item(a[i] + b[0], 0));
    }
    for(int i = 0; i < n; i++){
        Item tmp = pq.top();
        pq.pop();
        c[i] = tmp.s;
        int x = tmp.b;
        if(x + 1 < n){
            pq.push(Item(tmp.s - b[x] + b[x + 1], x + 1));
        }
    }
}

int main()
{
    int k;
    while(scanf("%d", &k) != EOF){
        for(int i = 0; i < k; i++){
            if(i == 0){
                for(int j = 0; j < k; j++){
                    scanf("%d", &a[j]);
                }
            }
            else {
                for(int j = 0; j < k; j++){
                    scanf("%d", &b[j]);
                }
                sort(b, b + k);
                merge(a, b, a, k);
            }
        }
        for(int i = 0; i < k; i++){
            if(i == k - 1) printf("%d\n", a[i]);
            else printf("%d ", a[i]);
        }
    }
    return 0;
}