上一篇文章我们用内置的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}]