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

Python 用 heapq 实战:Top-K、优先队列与合并多个有序流有个很常见的需求:从一百万条日志里找出耗时最长的 10 条。很多人第一反应是sorted(data, reverseTrue)[: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)elifxh[0]:# 比 K 个里最小的还大,才值得进来heapq.heapreplace(h,x)# 弹最小 压新值,一步完成比 poppush 快returnsorted(h,reverseTrue)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}]slowestheapq.nlargest(2,logs,keylambdar: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,taskheapq.heappop(pq)print(prio,task)# 1 付款 → 2 发邮件 → 3 归档这里有个真实会炸的坑:如果两个任务 priority 相同,元组会去比第二个元素task。task 要是个不可比较的对象(比如自定义类实例),直接TypeError: not supported。解决办法是插入一个自增序号做「平局裁判」,保证元组永远能比出大小,顺便实现了同优先级下的 FIFO:importheapq,itertools counteritertools.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()mergedheapq.merge(read_sorted(log1),read_sorted(log2),keylambdas:s[:19])小结heapq是标准库的最小堆,直接操作 list,堆顶h[0]永远是最小值。Top-K用大小为 K 的最小堆(或直接nlargest/nsmallest),O(n log k),省内存、能处理流式数据。优先队列用(priority, 序号, task)元组,加自增序号避免同优先级时比较到不可比对象。合并有序流用heapq.merge,惰性归并,内存只跟流数量相关。一句话记忆:要「全局有序」用 sort,要「局部最值/前几名/惰性归并」用 heap。分不清就问自己——我是不是根本不需要剩下那些数据的顺序?