python堆模块_python：用heapq内置模块实现堆

上一篇文章我们用内置的collections.deque双端队列实现了栈和FIFO队列，以及用maxlen参数实现了取最后的N个元素的操作，

但是除了简单队列，优先级队列---堆也是非常重要的一个数据结构，我们来看看如何用python实现吧。

heapq模块实现堆：

from heapq import heapify,heappush,heappop,nlargest,nsmallest

a = list(range(10))

heapify(a)   # 将堆属性应用到列表a上

heappush(a,-1)  # 入堆

heappop(a)     # 弹出并返回堆顶元素

nlargest(2,a)   #返回前n大的元素

nsmallest(2,a)  #返回前n小的元素

当然你可以自定义class，并实现__lt__()这个魔术方法，这样当你在堆中放这个类的对象时，堆会以你定义的对象的优先级返回元素：

或者在nsmallest和nlargest函数中传入key参数(这一点和sort一样)。

>>> import heapq

>>> portfolio = [

...     {'name': 'IBM', 'shares': 100, 'price': 91.1},

...     {'name': 'AAPL', 'shares': 50, 'price': 543.22},

...     {'name': 'FB', 'shares': 200, 'price': 21.09},

...     {'name': 'HPQ', 'shares': 35, 'price': 31.75},

...     {'name': 'YHOO', 'shares': 45, 'price': 16.35},

...     {'name': 'ACME', 'shares': 75, 'price': 115.65}

... ]

>>> cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])

>>> expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])

>>> cheap

[{'name': 'YHOO', 'shares': 45, 'price': 16.35}, {'name': 'FB', 'shares': 200, 'price': 21.09}, {'name': 'HPQ', 'shares': 35, 'price': 31.75}]

>>> expensive

[{'name': 'AAPL', 'shares': 50, 'price': 543.22}, {'name': 'ACME', 'shares': 75, 'price': 115.65}, {'name': 'IBM', 'shares': 100, 'price': 91.1}]