UVA 11997：K 個數組取前 K 小和（優先隊列歸併）

問題描述

給定 k 個長度均為 k 的數組，每個數組取一個數相加，共有 k^k 種和。求最小的 k 個和並輸出。

解題思路

利用「兩兩歸併」策略：先將第一個數組與第二個數組歸併得到前 k 小和，結果作為新數組再與第三個數組歸併，以此類推。兩個已排序數組求前 k 小和，可用優先隊列（最小堆）實現：初始將 A[i] + B[0] 入堆；每次取堆頂 (sum, idxB)，輸出 sum，若 idxB+1 < k，則推入 sum - B[idxB] + B[idxB+1]。

代碼實現 (C++)

#include 
using namespace std;
const int MAXK = 1000;
int A[MAXK][MAXK];
int n;

struct Node {
    int sum, idxB;
    bool operator<(const Node& o) const { return sum > o.sum; } // 最小堆
};

void mergeArrays(int* a, int* b, int* c) {
    priority_queue pq;
    for (int i = 0; i < n; ++i) pq.push({a[i] + b[0], 0});
    for (int i = 0; i < n; ++i) {
        Node cur = pq.top(); pq.pop();
        c[i] = cur.sum;
        if (cur.idxB + 1 < n) {
            pq.push({cur.sum - b[cur.idxB] + b[cur.idxB + 1], cur.idxB + 1});
        }
    }
}

int main() {
    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 = 1; i < n; ++i) mergeArrays(A[0], A[i], A[0]);
        for (int i = 0; i < n; ++i) {
            if (i) printf(" ");
            printf("%d", A[0][i]);
        }
        puts("");
    }
    return 0;
}

複雜度分析

• 排序：O(k * k log k) = O(k^2 log k)。
• 歸併：進行 k-1 次歸併，每次 O(k log k)，總 O(k^2 log k)。
• 空間：O(k^2) 存儲數組，O(k) 堆空間。

總結

本題將高維組合問題降維為反覆的二維歸併，利用優先隊列高效求兩數組前 k 小和，是經典的「K 最小和」問題模板。