python heapq

# 聊聊Python的heapq一个被低估的堆实现heapq是Python标准库里的一个模块它实现了堆队列算法也就是通常说的优先队列。这个模块不太起眼许多开发者甚至都不知道它的存在但它确实解决了不少让人头痛的问题。它到底是什么简单说heapq维护了一个列表让这个列表始终保持“堆”的特性。这个特性具体是列表的第一个元素总是最小的那个。如果往列表里添加新元素或者移除最小的那个元素它能在O(log n)的时间里完成调整保证你随时能从顶部拿到当前最小的元素。其实这就像你手头有一堆待办事项每次都挑最紧急的先处理。用heapq就能保证你每次都能快速找到那个“最紧急”的。它能做什么最直观的用途就是排序相关的场景。比如从海量数据里找出前K个最大的或最小的元素这就是典型的Top K问题。很多人可能第一反应是排序后取前K个但如果有10亿条数据呢排序再取前K个内存可能就爆了。还有一个场景是任务调度。假设你有一堆待处理的任务每个任务有不同的优先级你想每次取出优先级最高的那个来处理。用list然后每次找最大值复杂度是O(n)用heapq就是O(log n)差别挺大的。另外在网络爬虫的数据合并场景里也很有用。比如有多个已经排好序的数据源你想把它们归并成一个有序的整体heapq的归并功能就能帮上忙。怎么使用先看基本用法importheapq# 创建一个空堆heap[]# 添加元素heapq.heappush(heap,5)heapq.heappush(heap,2)heapq.heappush(heap,10)# 弹出最小元素smallestheapq.heappop(heap)# 返回2# 查看最小元素但不弹出peekheap[0]# 直接取索引0# 把列表转成堆data[5,2,10,1,8]heapq.heapify(data)# data现在变成了堆结构如果你处理的是复杂对象比如想按对象的某个属性排序可以传入一个元组importheapq heap[]heapq.heappush(heap,(priority,timestamp,data))这样heapq会先按元组的第一个元素比较如果相等再比较第二个以此类推。还有一个比较实用的函数是nlargest和nsmallestdata[5,2,10,1,8]largest_threeheapq.nlargest(3,data)# [10, 8, 5]smallest_threeheapq.nsmallest(3,data)# [1, 2, 5]一些实践中的考虑在使用heapq的时候有个容易踩的坑heapq默认实现的是最小堆。如果你想用最大堆可以存负值。比如想从一堆数里挑最大的就在入堆的时候存负数弹出的时候再取反。另外heapq的堆操作是原地修改列表的不会返回新的列表。这点和很多Python函数的行为不太一样容易让人意外。比如heapq.heappush(heap, item)不会返回新堆而是修改你传入的列表。如果你需要频繁地更新已有元素的值标准的heapq有点力不从心。比如任务调度中任务被插队或者优先级变更标准的堆实现不支持直接更新某个元素。这时候要么自己实现一个带索引的堆要么把更新操作转换成“新增一个元素”和“标记旧元素为无效”两个操作。后一种方法虽然内存占用稍大但实现起来简单得多。性能方面heapq对于百万级别的数据量表现很好。如果数据量大到内存装不下那就要考虑用外部排序或者数据库了。和其他方案的比较说到排序和选Top K首先能想到的是sorted函数。sorted会一次性把所有数据排序时间复杂度O(n log n)。如果用排序取Top K复杂度一样。而heapq取Top K只需要O(n log K)当K远小于n时优势明显。还有人会用列表的list.sort()方法效果和sorted类似区别在于就地排序。如果是单纯地每次取最大最小值用max/min也是O(n)而heapq是O(log n)。但如果只需要取一次用max/min就够了没必要建堆。另一个常被提到的数据结构是PriorityQueue它是线程安全的封装了heapq。如果你在多线程环境下使用PriorityQueue更省心。单线程下直接用heapq就好少一层包装性能更好。用第三方库的话有heapdict支持按键值更新堆元素还有sortedcontainers这种相对重量级的库。如果只是基本需求用heapq就够了没必要引入额外依赖。总体来说heapq是个小巧实用的工具适合那些需要频繁获取最小或最大元素但又不想每次都全排序的场景。它的API设计很Pythonic容易上手该有的功能都有不该有的也几乎寸没有。