python的heapq是一个优先队列的算法模块,它只提供函数算法,方便得到一个小根堆。
import heapq
s = [3,1,4,2,6]
heapq.heapify(s) #=> s = [1, 2, 4, 3, 6]
heapq.heappop(s) #=> 1
heapq.heappush(s, 0) #=> s = [0, 2, 4, 6, 3]
如果想最大的排前面,即一个大根堆,heapq不能直接支持设置比较函数,需要对数据进行简单处理:
ns = [(-i, i) for i in s]
heapq.heapify(ns)
heapq.heappop(ns)[1] #=> 6
小根堆在求取一个序列中最大的前N个数,或者最小的前N个数,比全排序然后截取的做法,效率要高。