标准库heapq的使用

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

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

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

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。

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

实现一个优先级队列

怎么实现一个按优先级排序的队列?并且在这个队列上面每次执行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()操作返回优先级最高的元素。另外注意到,有两个优先级一样的元素(defjkl),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做元组比较时,如果前面的比较可以确定了,后面的操作就不会发生了。

posted @ 2018-08-29 09:10  梁颖666  阅读(393)  评论(0编辑  收藏  举报