14 - heapq模块:堆操作

想系统提升编程能力、查看更完整的学习路线，欢迎访问 AI Compass：https://github.com/tingaicompass/AI-Compass

仓库持续更新刷题题解、Python 基础和 AI 实战内容，适合想高效进阶的你。

14 - heapq模块:堆操作

学习目标: 掌握Python的堆(优先队列)操作

💻 代码示例

python 复制代码

import heapq

# Python的堆是最小堆
nums = [3, 1, 4, 1, 5, 9, 2, 6]

# 将列表转换为堆
heapq.heapify(nums)
print(nums)  # [1, 1, 2, 3, 5, 9, 4, 6]

# 弹出最小元素
min_val = heapq.heappop(nums)
print(min_val)  # 1

# 添加元素
heapq.heappush(nums, 0)
print(heapq.heappop(nums))  # 0

# 替换堆顶
heapq.heapreplace(nums, 10)  # 弹出最小值,然后添加10

# 获取最大/最小的K个元素
nums = [3, 1, 4, 1, 5, 9, 2, 6]
print(heapq.nlargest(3, nums))   # [9, 6, 5]
print(heapq.nsmallest(3, nums))  # [1, 1, 2]

# 最大堆(取负数)
max_heap = []
heapq.heappush(max_heap, -5)
heapq.heappush(max_heap, -3)
heapq.heappush(max_heap, -8)
print(-heapq.heappop(max_heap))  # 8 (最大值)

🎯 在算法题中的应用

python 复制代码

# 第97课:前K个高频元素
import heapq
from collections import Counter

def topKFrequent(nums, k):
    count = Counter(nums)
    return heapq.nlargest(k, count.keys(), key=count.get)

# 第98课:数据流中位数(对顶堆)
class MedianFinder:
    def __init__(self):
        self.small = []  # 最大堆(存较小的一半)
        self.large = []  # 最小堆(存较大的一半)

    def addNum(self, num):
        if len(self.small) == len(self.large):
            heapq.heappush(self.large, -heapq.heappushpop(self.small, -num))
        else:
            heapq.heappush(self.small, -heapq.heappushpop(self.large, num))

    def findMedian(self):
        if len(self.small) == len(self.large):
            return (self.large[0] - self.small[0]) / 2
        else:
            return self.large[0]

# 第99课:合并K个升序链表
def mergeKLists(lists):
    import heapq
    heap = []

    # 初始化堆
    for i, node in enumerate(lists):
        if node:
            heapq.heappush(heap, (node.val, i, node))

    dummy = ListNode(0)
    current = dummy

    while heap:
        val, i, node = heapq.heappop(heap)
        current.next = node
        current = current.next

        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))

    return dummy.next

🎓 小结

✅ Python的堆是最小堆

✅ heapify(list) 将列表转换为堆

✅ heappush(heap, item) 添加元素

✅ heappop(heap) 弹出最小元素

✅ nlargest(k, iterable) 最大的K个

✅ nsmallest(k, iterable) 最小的K个

实现最大堆: 存储负数

下一步 : 16-高级特性.md

如果这篇内容对你有帮助，推荐收藏 AI Compass：https://github.com/tingaicompass/AI-Compass

更多系统化题解、编程基础和 AI 学习资料都在这里，后续复习和拓展会更省时间。