Python `heapq` 模块：优先队列与 K 个最大/最小元素问题

好的，各位听众，今天咱们来聊聊 Python 的 heapq 模块，以及它在优先队列和 K 个最大/最小元素问题上的应用。说白了，就是要教大家怎么用 Python 优雅地偷懒，高效地解决问题。

开场白：为什么我们需要 heapq？

想象一下，你是一家医院的急诊科医生，病人源源不断涌入，你得先救谁？肯定是最紧急的那个！这时候，你就需要一个“优先级”的概念。在计算机世界里，也经常遇到类似的情况。我们需要一种数据结构，能够快速找到“最重要”的元素，而不是每次都遍历一遍。

heapq 模块就是来解决这个问题的。它提供了一种基于堆（heap）的实现，可以高效地维护一个有序的集合，并且快速找到最大或最小的元素。这玩意儿就像一个自动排序的篮子，你往里面扔东西，它会自动把最重要的放在最前面。

什么是堆（Heap）？别害怕，其实很简单！

堆是一种特殊的树形数据结构，它满足以下两个性质：

1. 堆的形状： 堆是一个完全二叉树（除了最底层，其他层都是满的，最底层从左到右填充）。

2. 堆的性质： 堆中某个节点的值总是不大于（或不小于）其父节点的值。

   • 最小堆（Min-Heap）： 父节点的值小于或等于子节点的值。根节点是最小值。
   • 最大堆（Max-Heap）： 父节点的值大于或等于子节点的值。根节点是最大值。

heapq 模块默认实现的是最小堆。如果你需要最大堆，后面会告诉你怎么变通。

heapq 模块常用函数：你的工具箱

heapq 模块提供了一系列函数，让你轻松操作堆。下面是几个常用的：

函数	描述
heapq.heappush(heap, item)	将 item 放入 heap 中，保持堆的性质。
heapq.heappop(heap)	弹出并返回 heap 中最小的元素，同时保持堆的性质。
heapq.heapify(list)	将列表 list 转换为堆，原地修改列表。
heapq.heapreplace(heap, item)	弹出并返回 heap 中最小的元素，然后将 item 放入 heap 中，保持堆的性质。 效率比先 heappop 再 heappush 高。
heapq.nlargest(n, iterable, key=None)	返回 iterable 中最大的 n 个元素，以列表形式返回。可以指定 key 函数进行比较。
heapq.nsmallest(n, iterable, key=None)	返回 iterable 中最小的 n 个元素，以列表形式返回。可以指定 key 函数进行比较。

动手实践：代码才是硬道理

光说不练假把式，咱们来写点代码。

import heapq

# 创建一个空堆
my_heap = []

# 往堆里添加元素
heapq.heappush(my_heap, 5)
heapq.heappush(my_heap, 1)
heapq.heappush(my_heap, 3)
heapq.heappush(my_heap, 2)
heapq.heappush(my_heap, 4)

print(f"现在的堆：{my_heap}")  # 输出：现在的堆：[1, 2, 3, 5, 4]
# 注意：虽然列表看起来不是完全有序的，但它仍然满足堆的性质。

# 弹出最小的元素
smallest = heapq.heappop(my_heap)
print(f"弹出的最小元素：{smallest}")  # 输出：弹出的最小元素：1
print(f"现在的堆：{my_heap}")  # 输出：现在的堆：[2, 4, 3, 5]

# 从列表创建堆
my_list = [5, 1, 3, 2, 4]
heapq.heapify(my_list)
print(f"堆化后的列表：{my_list}")  # 输出：堆化后的列表：[1, 2, 3, 5, 4]

# 替换堆顶元素
heapq.heapreplace(my_list, 0)
print(f"替换后的堆：{my_list}") # 输出：替换后的堆：[0, 2, 3, 5, 4]

# 找到最大的 3 个元素
largest_3 = heapq.nlargest(3, [5, 1, 3, 2, 4])
print(f"最大的 3 个元素：{largest_3}")  # 输出：最大的 3 个元素：[5, 4, 3]

# 找到最小的 3 个元素
smallest_3 = heapq.nsmallest(3, [5, 1, 3, 2, 4])
print(f"最小的 3 个元素：{smallest_3}")  # 输出：最小的 3 个元素：[1, 2, 3]

优先队列：让任务排队，按重要性执行

优先队列是一种抽象数据类型，它允许你添加带有优先级的元素，并且总是先处理优先级最高的元素。heapq 模块可以轻松实现优先队列。

import heapq

class PriorityQueue:
    def __init__(self):
        self._queue = []
        self._index = 0  # 用于打破优先级相同的元素的顺序

    def push(self, item, priority):
        # heapq 默认是最小堆，所以我们用负数来表示优先级，值越大优先级越高
        heapq.heappush(self._queue, (-priority, self._index, item))
        self._index += 1

    def pop(self):
        return heapq.heappop(self._queue)[-1] # 提取item

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

# 使用优先队列
pq = PriorityQueue()
pq.push("任务 A", 3)
pq.push("任务 B", 1)
pq.push("任务 C", 2)
pq.push("任务 D", 3) # 优先级相同，按照进入队列的顺序

while not pq.is_empty():
    item = pq.pop()
    print(f"处理任务：{item}")
# 输出：
# 处理任务：任务 A
# 处理任务：任务 D
# 处理任务：任务 C
# 处理任务：任务 B

K 个最大/最小元素问题：从海量数据中找到 top K

假设你有一个巨大的数据集，比如一个包含 1 亿个数字的文件。你想找到其中最大的 10 个数字。如果直接排序，那效率就太低了。heapq 模块可以帮你解决这个问题。

import heapq
import random

def find_k_largest(nums, k):
    """
    找到列表中最大的 k 个元素。
    """
    return heapq.nlargest(k, nums)

def find_k_smallest(nums, k):
    """
    找到列表中最小的 k 个元素。
    """
    return heapq.nsmallest(k, nums)

# 生成一个包含 1000 个随机数的列表
random_numbers = [random.randint(1, 10000) for _ in range(1000)]

# 找到最大的 10 个数字
largest_10 = find_k_largest(random_numbers, 10)
print(f"最大的 10 个数字：{largest_10}")

# 找到最小的 10 个数字
smallest_10 = find_k_smallest(random_numbers, 10)
print(f"最小的 10 个数字：{smallest_10}")

# 大文件处理示例，避免一次性加载全部数据到内存
def find_k_largest_from_file(filename, k):
    """
    从文件中读取数字，找到最大的 k 个元素。
    假设文件每行包含一个数字。
    """
    largest_k = [] #存储k个最大值
    with open(filename, 'r') as f:
        for line in f:
            try:
                num = int(line.strip()) #转换成整数，并去除空白
                if len(largest_k) < k:
                    heapq.heappush(largest_k, num)
                elif num > largest_k[0]:
                    heapq.heapreplace(largest_k, num)
            except ValueError:
                print(f"忽略无效行：{line.strip()}") #处理非数字行

    return sorted(largest_k, reverse=True) #返回排序后的结果

#创建一个包含随机数的文件
filename = "random_numbers.txt"
with open(filename, "w") as f:
    for _ in range(1000):
        f.write(str(random.randint(1, 10000)) + "n")

#找到文件中最大的10个数字
largest_10_from_file = find_k_largest_from_file(filename, 10)
print(f"从文件中读取的最大的 10 个数字：{largest_10_from_file}")

进阶技巧：自定义比较规则

有时候，你需要的不仅仅是简单的数值比较，而是更复杂的比较规则。比如，你可能想根据对象的某个属性来排序。heapq 模块支持自定义比较函数。

import heapq

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

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

    def __lt__(self, other): #定义小于关系
        return self.priority < other.priority

tasks = [
    Task("任务 A", 3),
    Task("任务 B", 1),
    Task("任务 C", 2),
]

heapq.heapify(tasks)  #利用Task的__lt__方法构建堆

while tasks:
    task = heapq.heappop(tasks)
    print(f"处理任务：{task}")

#输出：
#处理任务：Task(name='任务 B', priority=1)
#处理任务：Task(name='任务 C', priority=2)
#处理任务：Task(name='任务 A', priority=3)

最大堆的实现：小技巧，大用处

heapq 模块默认是最小堆，但如果你需要最大堆，可以稍微变通一下：

1. 取反： 将所有元素取反，然后放入堆中。取出元素时，再取反回来。
2. 自定义比较函数： 使用 key 参数，自定义比较函数，实现最大堆的逻辑。

import heapq

# 使用取反实现最大堆
my_list = [5, 1, 3, 2, 4]
negated_list = [-x for x in my_list]  # 取反
heapq.heapify(negated_list)

largest = -heapq.heappop(negated_list)  # 取出时再取反
print(f"最大的元素：{largest}")

# 使用自定义比较函数实现最大堆（更推荐，避免类型问题）
class MaxHeap:
    def __init__(self, data=None):
        self._data = [-x for x in (data or [])]
        heapq.heapify(self._data)

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

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

    def peek(self):
        return -self._data[0]

    def __len__(self):
        return len(self._data)

max_heap = MaxHeap([5, 1, 3, 2, 4])
print(f"最大堆堆顶元素：{max_heap.peek()}") # 输出: 5
print(f"弹出最大元素：{max_heap.pop()}") # 输出: 5

性能分析：heapq 到底有多快？

heapq 模块的性能非常出色。

• heappush 和 heappop 的时间复杂度都是 O(log n)，其中 n 是堆的大小。
• heapify 的时间复杂度是 O(n)。
• nlargest 和 nsmallest 的时间复杂度是 O(n log k)，其中 n 是可迭代对象的长度，k 是要查找的元素个数。

这意味着，即使处理海量数据，heapq 也能保持高效。

应用场景：heapq 无处不在

heapq 模块的应用非常广泛，以下是一些常见的场景：

• 任务调度： 优先队列可以用于任务调度，确保优先级高的任务先执行。
• 图算法： Dijkstra 算法和 A* 算法等图算法中，需要使用优先队列来选择下一个要访问的节点。
• 数据流处理： 在数据流处理中，可以使用 heapq 来维护滑动窗口中的最大或最小元素。
• Top K 问题： 找到海量数据中最大的或最小的 K 个元素。
• 合并排序列表： 将多个已排序的列表合并成一个有序列表。

总结：heapq，你的得力助手

heapq 模块是 Python 标准库中一个非常实用的工具。它提供了一种高效的方式来处理优先队列和 K 个最大/最小元素问题。掌握 heapq 模块，可以让你在解决实际问题时更加得心应手。

记住，编程的本质是解决问题。heapq 模块就像一个强大的武器，可以帮助你更轻松地战胜各种难题。希望今天的讲解能让你对 heapq 模块有更深入的了解，并在实际应用中发挥它的威力。

课后作业：

1. 用 heapq 模块实现一个简单的任务管理器，可以添加任务，设置优先级，并按照优先级执行任务。
2. 从一个包含 100 万个随机数的文件中，找到最大的 100 个数字，并计算它们的平均值。

祝大家编程愉快！下课！