标准库heapq的使用

转载自： https://blog.csdn.net/y472360651/article/details/80725355

查找最大或最小的N个元素

怎么样从一个列表中取出最大或最小的N个元素的列表？在Python的标准库中，有一个名为heapq的，该模块中具有两个函数nlargest和nsmallest可以完全解决我们的问题，下面我们来看看这两个函数的作用：

import heapq

L = [5, 4, 6, 2, 8, 10, 1]

# 获取列表中最大的三个元素
print heapq.nlargest(3, L)
# 获取列表中最小的三个元素
print heapq.nsmallest(3, L)

# ===========打印结果如下===========
[10, 8, 6]
[1, 2, 4]

两个函数都可以传递参数，用于更复杂的数据结构当中：

import heapq

info_list = [
    {"name": "laozhang", "age": 20, "score": 90},
    {"name": "laoli", "age": 21, "score": 88},
    {"name": "laowang", "age": 24, "score": 58},
    {"name": "laohe", "age": 22, "score": 77},
    {"name": "laoyang", "age": 21, "score": 89}
]

# 取出年龄最大的两个数据
print heapq.nlargest(2, info_list, key=lambda item: item["age"])

# ===========打印结果如下===========
[{'name': 'laowang', 'age': 24, 'score': 58}, {'name': 'laohe', 'age': 22, 'score': 77}]

其实，在底层的实现中，heapq首先会将列表数据进行堆排序后放入一个列表：

import heapq

L = [5, 4, 6, 8, 2, 7]
# heapify操作之后，会将列表中最小的元素挪到第一位
heapq.heapify(L)
print L                 # [2, 4, 6, 8, 5, 7]

print heapq.heappop(L)  # 2
print heapq.heappop(L)  # 4
print heapq.heappop(L)  # 5

堆数据结构最重要的特征就是，列表的第一个元素永远是最小值。剩余的小值元素可以通过heapq.heappop(iterable)获得。该方法会先将第一个元素弹出来，然后下一个最小的元素来取代被弹出的元素，以此来获得列表L中最小的三个元素为2，4，5。

当查找的元素个数相对比较少时，nlargest和nsmallest函数是很合适的。如果你仅仅想查找唯一一个最大的或最小的值时，那么min和max函数会更快一些。类似的，如果查找元素个数与列表的长度大小接近时，通常先排序这个列表，然后使用切片的方式会更加快一些，我们需要在正确的场合使用nlargest和nsmallest函数才能发挥它们的优势。

实现一个优先级队列

怎么实现一个按优先级排序的队列？并且在这个队列上面每次执行pop操作总是返回优先级最高的那个元素。下面的类使用heaqp模块实现了一个简单的优先级队列：

import heapq

class HeapQueue(object):
    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):
        return heapq.heappop(self._queue)

class Item(object):
    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return "Item({!r})".format(self.name)

q = HeapQueue()
q.push(Item("abc"), 5)
q.push(Item("def"), 1)
q.push(Item("ghi"), 2)
q.push(Item("jkl"), 1)

print q._queue
print q.pop()

print q._queue
print q.pop()

print q._queue
print q.pop()

print q._queue
print q.pop()

# =============打印结果如下=============
[(-5, 0, Item('abc')), (-1, 1, Item('def')), (-2, 2, Item('ghi')), (-1, 3, Item('jkl'))]
(-5, 0, Item('abc'))

[(-2, 2, Item('ghi')), (-1, 1, Item('def')), (-1, 3, Item('jkl'))]
(-2, 2, Item('ghi'))

[(-1, 1, Item('def')), (-1, 3, Item('jkl'))]
(-1, 1, Item('def'))

[(-1, 3, Item('jkl'))]
(-1, 3, Item('jkl'))

从打印结果观察可以得知，第一个pop()操作返回优先级最高的元素。另外注意到，有两个优先级一样的元素（def和jkl），pop操作是按照它们的插入顺序返回的，先插入的先返回，后插入的后返回。

在上面代码中，队列包含了一个(-priority, self._index, item)元组，优先级为负数的目的是使得元素按照优先级从高到低进行排序，index变量确保同等优先级下，按照插入顺序先后取出，先插入的先取出，后插入的后取出。

首先阐明一点，对象的实例是不能进行排序的，所以我们必须引入一个元组，并对其设置优先级，只要对两个元素的优先级不同就能进行比较，如果两个优先级一样，就需要引入一个index变量，作为辅助排序。如果没有index作辅助排序，那么两个优先级相同的就不能进行比较，将会出错，请看下面错误比较：

a = (1, Item("abc"))
b = (1, Item("def"))
print a < b

# 将会报如下错误：
TypeError: '<' not supported between instances of 'Item' and 'Item'

但是，如果我们引入一个变量index作辅助排序，就能进行比较：

a = (1, 3, Item("abc"))
b = (1, 1, Item("def"))
print(a < b)

#===========打印结果如下===========
False

因为在本示例中，index变量不可能会出现相同的值，所以就可以很好的避免上面的错误。在Python做元组比较时，如果前面的比较可以确定了，后面的操作就不会发生了。