heapq 模块

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 "", line 1, in 
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
>>>