Python 用 heapq 实战:Top-K、优先队列与合并多个有序流

Python 用 heapq 实战:Top-K、优先队列与合并多个有序流

有个很常见的需求:从一百万条日志里找出耗时最长的 10 条。很多人第一反应是sorted(data, reverse=True)[:10]——排完序取前 10。能跑,但你为了 10 条结果把一百万条全排了一遍,O(n log n)的功全花在你根本不要的那 99.9% 数据上。

这类「从大堆里挑最大/最小的几个」问题,正确工具是堆(heap)。Python 标准库的heapq就是一个基于列表的最小堆实现,不用装任何东西。这篇把它三个最实用的场景讲透。

先认识 heapq:它是「最小堆」

heapq直接操作普通 list,把它当堆用。核心两个操作:heappush压入、heappop弹出最小值。

importheapq h=[]heapq.heappush(h,5)heapq.heappush(h,1)heapq.heappush(h,3)print(heapq.heappop(h))# 1,永远先弹最小的print(heapq.heappop(h))# 3

记住一句话:堆顶(h[0])永远是当前最小值,heappop弹的也是它。这是后面所有技巧的基础。

场景一:Top-K 最大值,别全排序

要找最大的 K 个,直觉是「用最大堆」,但更省内存的做法是维护一个大小为 K 的最小堆:遍历数据,堆没满就压;满了就拿新值跟堆顶(当前 K 个里最小的)比,比它大就替换掉堆顶。

importheapqdeftop_k(nums,k):h=[]forxinnums:iflen(h)<k:heapq.heappush(h,x)elifx>h[0]:# 比 K 个里最小的还大,才值得进来heapq.heapreplace(h,x)# 弹最小 + 压新值,一步完成比 pop+push 快returnsorted(h,reverse=True)print(top_k([3,1,9,7,2,8,5],3))# [9, 8, 7]

复杂度O(n log k),内存只占 K 个元素。数据是流式来的(读不完的日志、网络包)时,这个「只留 K 个」的特性尤其关键——你没法先攒齐再排序。

标准库其实封装好了这个模式:heapq.nlargest(k, nums)和nsmallest(k, nums),内部就是上面的逻辑,还支持key:

logs=[{"url":"/a","ms":120},{"url":"/b","ms":998},{"url":"/c","ms":45}]slowest=heapq.nlargest(2,logs,key=lambdar:r["ms"])

一个选型提醒:K 很小(比如 top 10)用nlargest;但如果 K 接近 n(比如取前 90%),直接sorted反而更快,别硬套堆。

场景二:优先队列,用元组带 priority

任务调度里经常要「优先级高的先执行」。堆本身只会按元素大小排,想按优先级排,就把(priority, task)打包成元组压进去——元组比较是逐位比的,先比第一个元素。

importheapq pq=[]heapq.heappush(pq,(2,"发邮件"))heapq.heappush(pq,(1,"付款"))# 数字小 = 优先级高heapq.heappush(pq,(3,"归档"))whilepq:prio,task=heapq.heappop(pq)print(prio,task)# 1 付款 → 2 发邮件 → 3 归档

这里有个真实会炸的坑:如果两个任务 priority 相同,元组会去比第二个元素task。task 要是个不可比较的对象(比如自定义类实例),直接TypeError: '<' not supported。解决办法是插入一个自增序号做「平局裁判」,保证元组永远能比出大小,顺便实现了同优先级下的 FIFO:

importheapq,itertools counter=itertools.count()# 全局自增计数器pq=[]defpush(pq,priority,task):# (优先级, 序号, 任务):序号唯一,元组比较永远不会落到 task 上heapq.heappush(pq,(priority,next(counter),task))push(pq,1,object())# 不可比较的对象也没问题push(pq,1,object())# 同优先级,按入队顺序出

场景三:合并 K 个有序流,不爆内存

假设你有 10 个已经按时间排好序的日志文件,要合成一条全局有序的流。全读进内存再sorted是最蠢的做法——文件可能几个 G。heapq.merge专为此而生,它接收多个有序可迭代对象,惰性地产出一个全局有序的迭代器:

importheapq a=[1,4,7]b=[2,5,8]c=[3,6,9]forxinheapq.merge(a,b,c):print(x,end=" ")# 1 2 3 4 5 6 7 8 9

关键在「惰性」:merge任意时刻只在堆里放 K 个元素(每个流的当前头),内存占用跟流的数量有关,跟总数据量无关。配合生成器读文件,几个 G 的日志归并也只占几 KB 内存:

defread_sorted(path):withopen(path)asf:forlineinf:# 逐行 yield,不一次性读入yieldline.rstrip()merged=heapq.merge(read_sorted("log1"),read_sorted("log2"),key=lambdas:s[:19])

小结

• heapq是标准库的最小堆,直接操作 list,堆顶h[0]永远是最小值。
• Top-K用大小为 K 的最小堆(或直接nlargest/nsmallest),O(n log k),省内存、能处理流式数据。
• 优先队列用(priority, 序号, task)元组,加自增序号避免同优先级时比较到不可比对象。
• 合并有序流用heapq.merge,惰性归并,内存只跟流数量相关。

一句话记忆:要「全局有序」用 sort,要「局部最值/前几名/惰性归并」用 heap。分不清就问自己——我是不是根本不需要剩下那些数据的顺序?