LeetCode 补充题 100. 合并 k 个有序数组

约 5 分钟
文章信息 6 个分类 1 个标签
分类
中等 已掌握 历史考题 补充题 字节 滴滴
标签
归并排序

题目描述

给你 k 个按非降序排列的整数数组 arrays，请将它们合并成一个按非降序排列的新数组，并返回结果。

各数组的长度可以不同，允许存在空数组和重复元素。结果需要保留所有元素及其出现次数。

示例 1：

输入： arrays = [[1,4],[],[1,3]]
输出： [1,1,3,4]
解释： 空数组不提供元素，重复出现的 1 都需要保留。

示例 2：

输入： arrays = [[],[]]
输出： []

提示：

  • k == arrays.length。
  • 每个数组均按非降序排列，数组长度可以不同。
  • 允许存在空数组和重复元素。

题意分析

直接合并所有元素后排序，时间复杂度为 $O(T \log T)$，其中 T 是元素总数，但没有利用每个数组已经有序的条件。

合并结果的下一个元素，一定是各数组“当前未合并的第一个元素”中的最小值。因此可以用最小堆维护这 k 个候选值。

堆节点除了保存元素值，还要保存它所在的数组下标和元素下标，这样弹出堆顶后，才能将同一数组的下一个元素加入堆。

解法：最小堆 + 多路归并

核心思路

[!blue]

堆中的每一项保存 value、row、col，分别表示候选值、来源数组和在该数组中的位置。每个尚未取完的非空数组恰好贡献一个候选，即它当前未输出的第一个元素。

因为各数组已经有序，同一数组后面的元素都不小于这个候选，所以所有剩余元素的最小值必然在堆顶。弹出它作为下一项后，只补入同一数组的后继，便恢复了这个不变量。

空数组不提供候选；某一路耗尽后不再入堆。相同数值来自不同位置时仍是独立元素，都会输出。所有候选消耗完时，结果既有序，也恰好包含全部输入元素。

解题步骤

  1. 检查所有数组，将非空数组的首元素及其来源位置加入最小堆。
  2. 弹出堆顶，将其值追加到结果末尾。
  3. 根据 row、col 找到同一数组的下一项，存在时加入堆。
  4. 重复直到堆为空；Java 预先统计总长度分配结果，Go 逐项追加结果。

代码实现

class Solution {
    public int[] mergeKArrays(int[][] arrays) {
        // 堆节点：{元素值, 数组下标, 元素下标}
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> Integer.compare(a[0], b[0]));
        int total = 0;

        for (int i = 0; i < arrays.length; i++) {
            total += arrays[i].length;

            if (arrays[i].length > 0) {
                heap.offer(new int[] {
                    arrays[i][0],
                    i,
                    0
                });
            }
        }

        int[] result = new int[total];
        int index = 0;

        while (!heap.isEmpty()) {
            // 堆中每个数组只提供当前最小未取元素，堆顶就是下一输出值。
            int[] current = heap.poll();

            result[index++] = current[0];

            int row = current[1];
            // 只补入刚弹出元素所属数组的下一个候选。
            int col = current[2] + 1;

            if (col < arrays[row].length) {
                heap.offer(new int[] {
                    arrays[row][col],
                    row,
                    col
                });
            }
        }

        return result;
    }
}

import "container/heap"

type item struct {
    value, row, col int
}

type minHeap []item

func (h minHeap) Len() int {
    return len(h)
}

func (h minHeap) Less(i, j int) bool {
    return h[i].value < h[j].value
}

func (h minHeap) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
}

func (h *minHeap) Push(x any) {
    *h = append(*h, x.(item))
}

func (h *minHeap) Pop() any {
    old := *h
    x := old[len(old)-1]
    *h = old[:len(old)-1]
    return x
}

func mergeKArrays(arrays [][]int) []int {
    h := &minHeap{}
    result := make([]int, 0)

    for row, array := range arrays {
        if len(array) > 0 {
            heap.Push(h, item{array[0], row, 0})
        }
    }

    for h.Len() > 0 {
        // 堆中每个数组只提供当前最小未取元素，堆顶就是下一输出值。
        current := heap.Pop(h).(item)
        result = append(result, current.value)

        // 只补入刚弹出元素所属数组的下一个候选。
        row, col := current.row, current.col+1
        if col < len(arrays[row]) {
            heap.Push(h, item{arrays[row][col], row, col})
        }
    }

    return result
}

复杂度分析

  • 时间复杂度：$O(k+T\log(k+1))$，先检查 k 个数组，再逐个处理总计 T 个元素；空数组也需要检查，堆大小最多为 k。
  • 空间复杂度：$O(k)$，不计输出数组。

关键点总结

[!green]

  • 堆中只保留每个数组当前的候选元素，堆大小不超过 k。
  • 每次弹出最小值后，只需补入同一数组的下一个元素。
  • 堆节点必须保存数组下标和元素下标。
  • 空数组不入堆，重复元素可以正常合并。

易错点总结

[!yellow]

  • 直接访问空数组的第一个元素，导致越界。
  • 弹出堆顶后忘记补入同一数组的下一个元素。
  • 堆节点只保存值，无法确定下一个元素来自哪个数组。
  • Java 堆比较器直接使用 a[0] - b[0]，整数溢出时比较结果会出错。

相似题目

题目 难度 关联与区别
23. 合并 K 个升序链表 困难 同样维护 k 路当前最小候选，本题通过数组下标取后继，原题沿链表 next 取后继。
373. 查找和最小的 K 对数字 中等 同样从多个有序来源逐次取最小项，数对和题需要额外定义每一路候选的生成方式。
转载与许可
作者
链接 https://hgnulb.github.io/blog/2025/77650347
许可 本博客所有文章除特别声明外，均采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处！