跳转至

优先队列

普通队列按照进入顺序取出元素;优先队列按照优先级取出元素。它适合“每次选出当前最小或最大元素,同时持续加入新元素”的场景,例如合并多个有序序列。

最小堆与基本操作

Python 的 heapq 把列表视为最小堆,堆顶是最小元素。堆不是完全排序的列表,除了堆顶外,不能假设任意位置已经有序。

import heapq

tasks = [7, 2, 5]
heapq.heapify(tasks)
heapq.heappush(tasks, 1)
assert heapq.heappop(tasks) == 1
assert tasks[0] == 2
操作 时间复杂度
从列表建堆 heapify O(n)
查看堆顶 heap[0] O(1)
入堆、出堆 O(log n)
逐个取出全部元素 O(n log n)

相同优先级怎样处理?

任务对象不一定可比较。使用 (优先级, 序号, 内容),让唯一递增的序号在优先级相同时决定先后顺序,也避免比较内容对象。

import heapq
from itertools import count

queue = []
sequence = count()

def push_task(priority, task):
    heapq.heappush(queue, (priority, next(sequence), task))

push_task(2, {"name": "整理笔记"})
push_task(1, {"name": "修复错误"})
push_task(2, {"name": "补充例子"})

assert heapq.heappop(queue)[2]["name"] == "修复错误"
assert heapq.heappop(queue)[2]["name"] == "整理笔记"

这里约定数字越小越优先。序号让相同优先级的任务按加入顺序取出。

用最小堆维护最大的 k 个数

堆中只保存当前最大的 k 个数。堆顶是这些候选中最小的值;新元素比堆顶更大时,替换堆顶。

import heapq

def top_k_largest(values, k):
    if k <= 0:
        return []
    heap = []
    for value in values:
        if len(heap) < k:
            heapq.heappush(heap, value)
        elif value > heap[0]:
            heapq.heapreplace(heap, value)
    return sorted(heap, reverse=True)

assert top_k_largest([3, 1, 5, 5, 2], 3) == [5, 5, 3]
assert top_k_largest([], 3) == []
assert top_k_largest([1, 2], 0) == []

当 1 ≤ k ≤ n 时,遍历部分耗时 O(n log k),最终排序耗时 O(k log k),额外空间 O(k)。k 大于元素数量时返回全部元素;重复值保留。

不要直接修改堆中的优先级

原地修改元素可能破坏堆结构。需要更新优先级时,可以重新建堆,或者使用“旧条目标记失效、插入新条目”的方式,并在取出时跳过失效条目。

延伸练习

结合二叉堆理解结构,再尝试合并 k 个升序链表:堆中为什么只需要保存每条链表当前的头节点?