深入探索 Python 的 heapq 模块：实现高效的优先队列

Jr_l

于 2024-07-15 09:00:00 发布

阅读量615

点赞数 6

分类专栏： python 文章标签： python

本文链接：https://blog.csdn.net/LS_Ai/article/details/140424793

版权

python 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

在现代编程中，优先队列（Priority Queue）是一种重要的数据结构。它允许我们在一系列元素中快速找到最小或最大的元素，并进行插入和删除操作。Python 提供了强大的 heapq 模块，使我们能够轻松实现这些功能。本文将深入探讨 heapq 模块的使用，并通过一个具体的例子展示其实际应用。

什么是堆？

堆是一种特殊的树形数据结构，可以被视为一个完全二叉树。堆分为两种类型：

最小堆：每个父节点的值都小于或等于其子节点的值。堆顶元素是最小值。
最大堆：每个父节点的值都大于或等于其子节点的值。堆顶元素是最大值。

在 Python 中，heapq 模块实现的是最小堆。如果需要最大堆，可以通过存储元素的负值来实现。

heapq 模块的常用函数

1. heapq.heappush(heap, item)

将元素 item 添加到堆 heap 中，并保持堆的特性。

import heapq

heap = []
heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
heapq.heappush(heap, 4)
print(heap)  # 输出: [1, 3, 4]

2. heapq.heappop(heap)

从堆中弹出并返回最小元素，同时保持堆的特性。

min_item = heapq.heappop(heap)
print(min_item)  # 输出: 1

3. heapq.heapify(iterable)

将一个可迭代对象转换为堆。

numbers = [4, 1, 7, 3, 8, 5]
heapq.heapify(numbers)
print(numbers)  # 输出: [1, 3, 5, 4, 8, 7]

4. heapq.nlargest(n, iterable, key=None)

返回可迭代对象中最大的 n 个元素。

largest = heapq.nlargest(3, numbers)
print(largest)  # 输出: [8, 7, 5]

5. heapq.nsmallest(n, iterable, key=None)

返回可迭代对象中最小的 n 个元素。

smallest = heapq.nsmallest(3, numbers)
print(smallest)  # 输出: [1, 3, 4]

示例应用：最后一块石头的重量

我们通过一个 LeetCode 问题来演示 heapq 模块的实际应用。问题描述如下：

有一堆石头，每块石头的重量都是正整数。每一回合，从中选出两块最重的石头，然后将它们一起粉碎。如果它们的重量相等，则两块石头都会完全粉碎；如果不等，则重量较大的石头会减小到原来的重量减去较小的石头的重量。最终，最多只会剩下一块石头。返回此石头的重量。如果没有石头剩下，则返回 0。

实现方案

我们将使用 heapq 模块来解决这个问题。由于 heapq 是最小堆，我们需要将所有石头的重量取反，从而使用最小堆模拟最大堆的行为。

import heapq
from typing import List

class Solution:
    def lastStoneWeight(self, stones: List[int]) -> int:
        heap = [-stone for stone in stones]  # 将所有石头重量取反
        heapq.heapify(heap)  # 将列表转换为堆

        while len(heap) > 1:
            x = heapq.heappop(heap)  # 弹出最重的石头（负值）
            y = heapq.heappop(heap)  # 弹出次重的石头（负值）
            if x != y:
                heapq.heappush(heap, x - y)  # 如果两块石头重量不等，将差值（仍为负值）压入堆

        return -heap[0] if heap else 0  # 如果堆不为空，返回堆中剩下的唯一元素的相反数

# 测试代码
stones = [2, 7, 4, 1, 8, 1]
solution = Solution()
print(f"Last stone weight: {solution.lastStoneWeight(stones)}")  # 输出: 1