heapq--堆的使用以及实现优先级队列

2019年5月19日 108次阅读 来源: cook__

抛出问题：找到最大或最小的N个元素
heapq模块中有两个函数nlargest()和nsmallest()，可以分别返回集合中最大和最小的元素。

import heapq
nums = [1, -2, 0, 34, 1, 6, -8]
heapq.nlargest(3, nums)
Out[4]: [34, 6, 1]
heapq.nsmallest(3, nums)
Out[5]: [-8, -2, 0]

这两个函数都可以接受一个参数key，从而允许他们工作在更复杂的数据结构之上。

portfolio = [
    {'name': 'IBM', 'shares': 100, 'prices': 91.1},
    {'name': 'IBM', 'shares': 100, 'prices': 91.2},
    {'name': 'IBM', 'shares': 100, 'prices': 91.4},
    {'name': 'IBM', 'shares': 100, 'prices': 91.9}
]
heapq.nlargest(3, portfolio, key=lambda s: s['prices'])
Out[7]: 
[{'name': 'IBM', 'prices': 91.9, 'shares': 100},
 {'name': 'IBM', 'prices': 91.4, 'shares': 100},
 {'name': 'IBM', 'prices': 91.2, 'shares': 100}]
heapq.nsmallest(3, portfolio, key=lambda s: s['prices'])
Out[8]: 
[{'name': 'IBM', 'prices': 91.1, 'shares': 100},
 {'name': 'IBM', 'prices': 91.2, 'shares': 100},
 {'name': 'IBM', 'prices': 91.4, 'shares': 100}]

如果正在寻找最大或最小的N个元素，且同集合中元素的总数相比，N很小，那么下面这些函数可以提供更好的性能。这些函数首先会在底层将数据转化成列表，且元素会以堆的顺序排列。

nums = [1, -2, 0, 34, 1, 6, -8, 12, -11, -9, 102]
heap = list(nums)
heapq.heapify(heap)
heap
Out[12]: [-11, -9, -8, -2, 1, 6, 0, 12, 34, 1, 102]
堆最重要的特性就是heap[0]总是那个最小的元素。

接下来的元素可依次通过heapq.heappop()方法找到，该方法会将第一个最小的元素弹出，然后以第二小的元素取而代之。
例如，要找到第三小的元素：

Out[12]: [-11, -9, -8, -2, 1, 6, 0, 12, 34, 1, 102]
heapq.heappop(heap)
Out[13]: -11
heapq.heappop(heap)
Out[14]: -9
heapq.heappop(heap)
Out[15]: -8

总结一下集合中查找最大或最小的常用方法：
1、当所要找的元素数量相对较小时，函数nlargest()和nsmallest()才是最适用的；
2、如果只是简单想找到最大或最小的元素(N=1时)，使用max()和min()会更加快；
3、如果N和集合本身的大小差不多大，通常更快的方法是先对集合排序，然后做切片操作(使用sorted(items)[:N]或者sorted(items)[-N:])；

在此，针对3做法补充一点：
list.sort() 会改变list对象，返回None；
sorted(list) 不改变原list对象，返回新生成的排序后的list对象；

堆的应用： –实现优先级队列

要求：实现一个队列，它能够以给定的优先级来对元素进行排序，且每次pop操作时都会返回优先级最高的那个元素。

class PriorityQueue:
    """ 用堆实现优先级队列 """
    def __init__(self):
        self._queue = []
        self._index = 0

    def push(self, item, priority):
        """ 将self._queue定义为heap堆, 将(-priority, self._index, item)定义为堆的元素 """
        # heapq.heappop()会取出最小元素，对堆元素(-priority, self._index, item))从左至右比较；因此_index能将具有相同优先级的元素按照插入顺序排列
        heapq.heappush(self._queue, (-priority, self._index, item))  # 插入元素
        self._index += 1

    def pop(self):
        """ 取出堆的item: heapq.heappop()方法一直弹出堆中最小的那个元素 """
        return heapq.heappop(self._queue)[-1]

使用示例：

class Item:
    def __init__(self, name):
        self.name = name
    def __repr__(self):
        return 'Item({!r})'.format(self.name)
    
q = PriorityQueue()
q.push(Item('foo'), 1)
q.push(Item('bar'), 5)
q.push(Item('spam'), 4)
q.push(Item('grok'), 1)
q.pop()
Out[10]: Item('bar')
q.pop()
Out[11]: Item('spam')
q.pop()
Out[12]: Item('foo')
q.pop()
Out[13]: Item('grok')
    原文作者：cook__
    原文地址: https://www.jianshu.com/p/5b83720e8da7
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞