问题描述

给定K个按升序排列的数组，要求将它们合并为一个大的有序数组。例如，输入数组[[1,3,5], [2,4,6], [0,7]]，合并后的结果应为[0,1,2,3,4,5,6,7]。

解决方案

思路分析

合并多个有序数组的高效方法是利用最小堆（优先队列）。核心步骤如下：

1. 初始化堆：将所有数组的第一个元素插入堆中，堆中每个元素记录其值、所属数组索引及元素在数组中的位置。

2. 循环处理堆顶元素：每次取出堆顶最小元素，将其加入结果数组，并将该元素所在数组的下一个元素（若存在）插入堆中。

3. 终止条件：当堆为空时，所有元素处理完毕。

此方法的时间复杂度为O(N log K)，其中N为总元素数，K为数组个数。每次堆操作的时间复杂度为O(log K)，总共有N次操作。

代码实现

#include <vector>
#include <queue>
using namespace std;

struct HeapNode {
    int val;
    int arrayIdx;
    int elementIdx;
    HeapNode(int v, int aIdx, int eIdx) : val(v), arrayIdx(aIdx), elementIdx(eIdx) {}
};

struct Compare {
    bool operator()(const HeapNode& a, const HeapNode& b) {
        return a.val > b.val; // 小顶堆
    }
};

vector<int> mergeKArrays(vector<vector<int>>& arrays) {
    vector<int> result;
    int k = arrays.size();
    if (k == 0) return result;

    priority_queue<HeapNode, vector<HeapNode>, Compare> minHeap;

    // 初始化堆，插入每个数组的第一个元素
    for (int i = 0; i < k; ++i) {
        if (!arrays[i].empty()) {
            minHeap.push(HeapNode(arrays[i][0], i, 0));
        }
    }

    while (!minHeap.empty()) {
        HeapNode node = minHeap.top();
        minHeap.pop();
        result.push_back(node.val);

        int nextElementIdx = node.elementIdx + 1;
        if (nextElementIdx < arrays[node.arrayIdx].size()) {
            minHeap.push(HeapNode(arrays[node.arrayIdx][nextElementIdx], node.arrayIdx, nextElementIdx));
        }
    }

    return result;
}

关键点解释

1. 结构体定义：HeapNode保存元素值、数组索引及元素位置，便于后续操作。

2. 比较函数：通过自定义比较结构体Compare，实现小顶堆，确保每次取出最小值。

3. 堆初始化：遍历所有数组，非空数组的首元素入堆。

4. 循环处理：取出堆顶元素后，若其所在数组还有后续元素，则将下一元素入堆。

边界条件处理

• 空数组：初始化时跳过空数组，避免无效操作。

• 全空输入：直接返回空结果。

• 单个数组：直接返回该数组本身。

测试案例

int main() {
    vector<vector<int>> test1 = {{1,3,5}, {2,4,6}, {0,7}};
    vector<int> res1 = mergeKArrays(test1);
    // 预期输出: 0 1 2 3 4 5 6 7

    vector<vector<int>> test2 = {{1,2}, {}, {3,4}};
    vector<int> res2 = mergeKArrays(test2);
    // 预期输出: 1 2 3 4

    return 0;
}

总结

通过最小堆高效管理各数组的当前最小元素，确保每次操作的时间复杂度为O(log K)，整体复杂度为O(N log K)。该方法直观且高效，适用于合并多个有序数据流的场景。