文章

算法刷题系列(二十五):多路归并——用小顶堆合并多个有序流

算法刷题系列(二十五):多路归并——用小顶堆合并多个有序流

算法刷题系列(二十五):多路归并——用小顶堆合并多个有序流

两个有序数组可以用双指针线性合并。若有 k 个有序数组,逐两两合并会重复搬运中间结果;更直接的办法是只看每个数组当前尚未处理的第一个元素,选择全局最小者输出,再推进它所在的数组。

这正好是最小堆的用途:堆中至多保留 k 个候选头部,每次弹出全局最小项。它与上一篇 Top-K 共用堆结构,但维护对象和目标不同——这里每个输入流都必须持续提供代表性候选。

一、归并步骤与例子

给定 [-2, 4, 10]、[1, 4]、[-1, 3, 12],初始堆为 -2、1、-1。先输出 -2 并从第一路补入 4;再输出 -1 并补入 3;之后每次输出堆顶、只推进该项所属的流。相等元素可以按任意流顺序输出,值的多重集合仍正确。

二、循环不变式

每轮开始时,堆中恰好有每个非空流的一个未输出头部。每条流内部已排序,因此该流未输出的其余元素都不小于当前头部。堆顶是所有未输出元素中的最小值,输出它不会破坏全局排序;将同一流的下一个元素加入后,不变式恢复。

三、Java 实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

public class Main {
    static final class Entry {
        final int value;
        final int listIndex;
        final int elementIndex;

        Entry(int value, int listIndex, int elementIndex) {
            this.value = value;
            this.listIndex = listIndex;
            this.elementIndex = elementIndex;
        }
    }

    static List<Integer> merge(int[][] lists) {
        PriorityQueue<Entry> heap =
                new PriorityQueue<>((a, b) -> Integer.compare(a.value, b.value));
        for (int i = 0; i < lists.length; i++) {
            if (lists[i].length > 0) {
                heap.offer(new Entry(lists[i][0], i, 0));
            }
        }

        List<Integer> result = new ArrayList<>();
        while (!heap.isEmpty()) {
            Entry current = heap.poll();
            result.add(current.value);
            int nextIndex = current.elementIndex + 1;
            if (nextIndex < lists[current.listIndex].length) {
                heap.offer(new Entry(
                        lists[current.listIndex][nextIndex],
                        current.listIndex,
                        nextIndex));
            }
        }
        return result;
    }

    public static void main(String[] args) {
        int[][] lists = {{-2, 4, 10}, {1, 4}, {-1, 3, 12}, {}};
        assert merge(lists).equals(List.of(-2, -1, 1, 3, 4, 4, 10, 12));
        assert merge(new int[][]{}).isEmpty();
        assert merge(new int[][]{{}, {}}).isEmpty();
    }
}

四、Python 实现

堆项使用 (值, 流编号, 下标) 元组。编号和下标也提供了稳定的次级比较键,避免值相等时尝试比较其他对象。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import heapq

def merge_sorted(lists: list[list[int]]) -> list[int]:
    heap: list[tuple[int, int, int]] = []
    for list_index, values in enumerate(lists):
        if values:
            heapq.heappush(heap, (values[0], list_index, 0))

    result: list[int] = []
    while heap:
        value, list_index, element_index = heapq.heappop(heap)
        result.append(value)
        next_index = element_index + 1
        if next_index < len(lists[list_index]):
            heapq.heappush(
                heap, (lists[list_index][next_index], list_index, next_index)
            )
    return result

assert merge_sorted([[-2, 4, 10], [1, 4], [-1, 3, 12], []]) == [
    -2, -1, 1, 3, 4, 4, 10, 12
]
assert merge_sorted([]) == []
assert merge_sorted([[], []]) == []

五、复杂度与边界

设总元素数为 N,非空流数为 k':

  • 时间:O(N log k');
  • 额外空间:O(k'),不包括结果数组。

空流不入堆,所有流为空时直接返回空结果。重复值保留其出现次数。算法要求每个输入流本身非递减;若输入并未排序,堆不能补救顺序错误。对外部文件或网络流,可只保留每路当前记录,因此内存仍与流数而不是总数据量成正比。

小结与衔接

多路归并的关键不是把所有数据一次性入堆,而是让每个有序来源只贡献一个当前候选。接下来从数据结构转向决策策略:贪心算法同样逐步选择,但其难点不是如何高效取最小值,而是如何证明局部选择能导向全局最优。

本文由作者按照 CC BY 4.0 进行授权