Python heapq 库详解：堆操作与实战应用

2.2 批量查询函数

代码示例：Top-K 问题解决

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']

2.3 高级操作函数

heapq.merge(*iterables) 函数用于合并多个已排序的输入序列，返回一个排序后的迭代器。

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]

3. 实战应用场景

3.1 优先级队列实现

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

3.2 实时数据流的中位数查找

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()}")

3.3 堆排序算法

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]

4. 高级技巧与性能优化

4.1 实现最大堆

由于 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

4.2 自定义对象堆操作

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))

5. 性能对比与最佳实践

5.1 不同场景下的性能选择

5.2 内存优化技巧

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}")

Python 的 heapq 库通过提供高效的堆操作函数，在算法优化、数据处理和系统设计等多个领域发挥着重要作用。掌握这些函数的正确使用方法和适用场景，能够显著提升程序的性能和代码的可维护性。