滴滴面试题-合并 k 个排序数组

2025-05-28 HGNULB 约 4 分钟

文章信息 5 个分类 1 个标签

分类

题目描述

给定 k 个升序数组，将它们合并成一个升序数组。

各数组的长度可以不同，允许存在空数组和重复元素。

题意分析

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

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

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

解法：最小堆 + 多路归并

核心思路

先将每个非空数组的第一个元素加入最小堆。

每次弹出堆顶加入结果，再将它所属数组的下一个元素入堆。整个过程中，每个未取完的数组在堆中最多只有一个元素，所以堆的大小不超过 k。

解题步骤

遍历所有数组，将每个非空数组的首元素加入最小堆。

弹出堆顶，将其值加入结果。

如果该元素所属数组还有下一个元素，将其加入堆。

重复上述过程，直到堆为空。

代码实现

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;
    }
}

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(T \log k)$，每个元素进堆、出堆各一次，堆大小最多为 k。

空间复杂度：$O(k)$，不计输出数组。

关键点总结

堆中只保留每个数组当前的候选元素，堆大小不超过 k。

每次弹出最小值后，只需补入同一数组的下一个元素。

堆节点必须保存数组下标和元素下标。

空数组不入堆，重复元素可以正常合并。

易错点总结

直接访问空数组的第一个元素，导致越界。

弹出堆顶后忘记补入同一数组的下一个元素。

堆节点只保存值，无法确定下一个元素来自哪个数组。

将复杂度写成 $O(T \log T)$；最小堆中最多只有 k 个元素。

Java 堆比较器直接使用 a[0] - b[0]，整数溢出时比较结果会出错。

相似题目

题目	难度	考察点
23. 合并 K 个升序链表	困难	最小堆多路归并
88. 合并两个有序数组	简单	双指针归并
373. 查找和最小的 K 对数字	中等	最小堆
378. 有序矩阵中第 K 小的元素	中等	最小堆多路归并