Python heapq 入门与实战：最小堆、Top-K 和优先队列

有个细节很容易被忽略：heapify() 只是把'堆性质'建立起来，不会帮你排成升序。很多人第一次看打印结果会误以为它是半排序状态，其实不是，堆只保证堆顶最小。

批量取前几个元素

import heapq
import random

# 生成测试数据
numbers = [random.randint(1, 1000) for _ in range(100)]

# 获取最大的 5 个元素
largest_5 = heapq.nlargest(5, numbers)
print(f"最大的 5 个元素：{largest_5}")

# 获取最小的 5 个元素
smallest_5 = heapq.nsmallest(5, numbers)
print(f"最小的 5 个元素：{smallest_5}")

# 使用 key 参数进行自定义比较
words = ['apple', 'banana', 'cherry', 'date', 'elderberry']
longest_3 = heapq.nlargest(3, words, key=len)
print(f"最长的 3 个单词：{longest_3}")
# 输出：['elderberry', 'banana', 'cherry']

nlargest() 和 nsmallest() 适合拿'前几名'，但别拿它去替代完整排序。数据量小的时候你感受不明显，数据一大，省下来的就是无意义的比较和内存开销。

合并多个有序序列

heapq.merge(*iterables) 会把多个已经排好序的序列合成一个有序迭代器。它不会一次性把所有数据拉进内存，这一点比先拼接再 sorted() 更省。

import heapq

# 合并多个有序序列
list1 = [1, 3, 5, 7]
list2 = [2, 4, 6, 8]
list3 = [0, 9, 10]
merged = list(heapq.merge(list1, list2, list3))
print(f"合并后的有序列表：{merged}")
# 输出：[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

优先级队列

heapq 最常见的实战用法之一就是优先级队列。Python 没有内置的稳定优先队列类型，但用一个三元组 (priority, index, item) 就够用了。index 的作用是打平相同优先级时的比较问题，也能顺便保留先来先服务的顺序。

import heapq

class PriorityQueue:
    def __init__(self):
        self._heap = []
        self._index = 0  # 用于处理优先级相同的情况

    def push(self, item, priority):
        """添加元素到优先级队列"""
        heapq.heappush(self._heap, (priority, self._index, item))
        self._index += 1

    def pop(self):
        """弹出优先级最高的元素"""
        if self._heap:
            return heapq.heappop(self._heap)[-1]
        raise IndexError("优先级队列为空")

    def is_empty(self):
        return len(self._heap) == 0

# 使用示例
pq = PriorityQueue()
pq.push("任务 A", 3)
pq.push("任务 B", 1)  # 最高优先级
pq.push("任务 C", 2)
while not pq.is_empty():
    print(f"执行：{pq.pop()}")
# 输出：执行：任务 B → 执行：任务 C → 执行：任务 A

数据流中位数

这个例子基本是堆的经典题：用一个最大堆存较小的一半，一个最小堆存较大的一半，随时保持两边平衡。heapq 没有真正的最大堆，只能把数取负来模拟。代码不算优雅，但在线求中位数时很好用。

import heapq

class MedianFinder:
    def __init__(self):
        # 最大堆（使用负数模拟）和最小堆
        self.max_heap = []  # 存储较小的一半
        self.min_heap = []  # 存储较大的一半

    def add_num(self, num):
        if not self.max_heap or num <= -self.max_heap[0]:
            heapq.heappush(self.max_heap, -num)
        else:
            heapq.heappush(self.min_heap, num)
        
        # 平衡两个堆
        if len(self.max_heap) > len(self.min_heap) + 1:
            heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
        elif len(self.min_heap) > len(self.max_heap):
            heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))

    def find_median(self):
        if len(self.max_heap) == len(self.min_heap):
            return (-self.max_heap[0] + self.min_heap[0]) / 2
        else:
            return -self.max_heap[0]

# 使用示例
finder = MedianFinder()
for num in [1, 3, 2, 6, 4, 5]:
    finder.add_num(num)
print(f"当前中位数：{finder.find_median()}")

堆排序

如果只是想验证堆的行为，堆排序也能写得很短。不过在 Python 里，真要排序通常直接用 sorted()，堆排序更多是教学用途，或者你想借堆结构顺手解决别的问题。

import heapq

def heap_sort(iterable):
    """使用堆排序算法对可迭代对象进行排序"""
    heap = list(iterable)
    heapq.heapify(heap)  # 构建最小堆
    return [heapq.heappop(heap) for _ in range(len(heap))]

# 排序示例
unsorted_data = [9, 2, 7, 5, 1, 8, 3, 6, 4]
sorted_data = heap_sort(unsorted_data)
print(f"堆排序结果：{sorted_data}")
# 输出：[1, 2, 3, 4, 5, 6, 7, 8, 9]

最大堆的写法

heapq 没有直接提供最大堆接口，常见做法就是把数取负。这个方案简单，够用，但只适合数值类型；如果是自定义对象，通常还是直接包装比较键更稳。

import heapq

class MaxHeap:
    def __init__(self):
        self._heap = []

    def push(self, item):
        heapq.heappush(self._heap, -item)

    def pop(self):
        return -heapq.heappop(self._heap)

    def peek(self):
        return -self._heap[0] if self._heap else None

# 最大堆使用示例
max_heap = MaxHeap()
for num in [3, 1, 4, 1, 5]:
    max_heap.push(num)
print("最大堆元素弹出顺序:")
while max_heap._heap:
    print(max_heap.pop())
# 输出：5, 4, 3, 1, 1

自定义对象的堆

自定义对象可以直接实现 __lt__，不过这类写法要小心：比较规则一旦写死，后面需求变了就容易牵一发动全身。只在对象本身就有明确排序语义时这么做，别为了省几行代码把逻辑绑死。

import heapq

class Task:
    def __init__(self, name, priority, duration):
        self.name = name
        self.priority = priority
        self.duration = duration

    def __lt__(self, other):
        # 定义比较规则：优先级高的在前，相同优先级时持续时间短的在前
        if self.priority == other.priority:
            return self.duration < other.duration
        return self.priority > other.priority

    def __repr__(self):
        return f"Task({self.name}, priority:{self.priority}, duration:{self.duration})"

# 自定义对象堆操作
tasks = [
    Task("紧急任务", 3, 2),
    Task("普通任务", 1, 5),
    Task("重要任务", 2, 3)
]
heap = []
for task in tasks:
    heapq.heappush(heap, task)
print("任务执行顺序:")
while heap:
    print(heapq.heappop(heap))

什么时候该用 heapq

如果你只是做一次完整排序，sorted() 通常更直接；如果你要的是'保持一个不断变化的候选集'，堆就更合适。它不花哨，但在合适的地方很好使。

大数据流里只保留 Top-N

import heapq

# 流式处理大数据集
def process_large_dataset(data_stream, top_n=10):
    """使用堆处理大数据流，只维护 Top-N 元素"""
    heap = []
    for item in data_stream:
        if len(heap) < top_n:
            heapq.heappush(heap, item)
        elif item > heap[0]:  # 对于最大 Top-N，使用最小堆
            heapq.heapreplace(heap, item)
    return sorted(heap, reverse=True)

# 模拟大数据流处理
import random
data_stream = (random.randint(1, 10000) for _ in range(100000))
top_10 = process_large_dataset(data_stream, 10)
print(f"大数据流中的 Top-10: {top_10}")

这一段的关键不在'快一点'，而在'别把所有数据都塞进内存'。只保留前 N 个元素，很多场景下就够了，尤其是日志、指标和排行榜一类的数据处理。

heapq 的使用门槛不高，真正容易出错的是选错结构。它适合维护局部最值、优先级和有序流，不适合拿来当通用排序工具。把这个边界记住，后面写起来会省很多事。