heapq 模块

标签（空格分隔）： pythoncook笔记

---

1. 查找最大或最小的N个元素（nlargest ,nsmallest）

import heapq nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2] print(heapq.nlargest(3, nums)) # Prints [42, 37, 23] print(heapq.nsmallest(3, nums)) # Prints [-4, 1, 2]

两个函数都能接受一个关键字参数，用于更复杂的数据结构中：

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'])

2. 优先队列的实现

import heapq class PriorityQueue: def __init__(self): self._queue = [] # 初始化一个列表来储存队列 self._index = 0 # 初始化信号 def push(self, item, priority): heapq.heappush(self._queue, (-priority, self._index, item)) # 向队列里添加元素，并添加优先级 self._index += 1 # 使每个添加进队列的元素具有不同的信号予以区分 def pop(self): # heapq.heappop会删除队列中最小的元素并返回，这正好照应了之前把优先级设为负数的构想 return heapq.heappop(self._queue)[-1]

index的作用是在两个元素具有相同的优先级的情况下，可以按照它们的插入顺序来将元素跳出队列 看一下代码你就知道为什么了

>>> a = (1, Item('foo')) >>> b = (5, Item('bar')) >>> a < b True >>> c = (1, Item('grok')) >>> a < c Traceback (most recent call last): File "<stdin>", line 1, in <module> TypeError: unorderable types: Item() < Item() >>>

当加上index变量组成三元组进行比较时，由于每个元素的index不一样，所以就能很好的避免上面的问题了

>>> a = (1, 0, Item('foo')) >>> b = (5, 1, Item('bar')) >>> c = (1, 2, Item('grok')) >>> a < b True >>> a < c True >>>