抛出问题:找到最大或最小的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')
网友评论