查找集合中N个最大或最小的元素-heapq
总结出以下3种解决问题的方案:
- 从集合中查找N个最大或最小的元素集,可将问题转化为每次从集合中查找最大或最小的元素,然后从集合中删除该元素。重复以上N次。
- 排序算法,按最大或最小顺序排序,然后从左至右选取集合中的N个元素。
- 使用堆序列算法(优先队列),构造二插堆,则集合中的首个元素即为最大值(俗称大堆),最小值(俗称小堆)。heap
本博文重点介绍优先队列的实现方案,并对比和排序算法直接的性能差异。
API
heapq模块直接提供了两个函数-nlargest()和nsmallest(),直接可以获取集合中最大或最小的N个元素集。
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]
堆是二叉树,每个父节点的值小于或等于其任何子节点。 对于所有k,此实现使用heap[k] <= heap [2 * k + 1]和heap [k] <= heap [2 * k + 2]的数组,从零开始计数元素。 为了比较,不存在的元素被认为是无限的。 堆的有趣属性是它的最小元素始终是根,heap[0]。这样,我们可以获取堆首元素,然后从集合中移除该元素,重新构造堆......
nlargest和nsmallest还可以接受一个关键字参数进行定制化比较,如:
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'])
如果要查找的N个最大或最小的元素和整个集合元素的数量相比很小,那么以上方法提供了非常好的性能。若N=1,即仅仅为了找到集合中最大或最小的单个元素,则最快的方式是使用max()或min()。若N~=len(collection),则最快的方式是使用排序算法,如sorted(items)[:N]。
实现优先队列
根据给定的优先级返回优先队列中的首个元素,可以按如下方式实现:
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):
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()
Item('bar')
>>> q.pop()
Item('spam')
>>> q.pop()
Item('foo')
>>> q.pop()
Item('grok')
>>>
引入_index可以确保若两个Item的优先级相同,且Item对象无法进行比较时,整个item的元素可以进行比较。
如果要在不同线程中使用queue,则需要添加相应的锁和信号(appropriate locking and signaling)。
heap API
class MinHeap:
def __init__(self, iterable=[]):
self.items = iterable
def _exch(self, i, j):
self.items[i], self.items[j] = self.items[j], self.items[i]
def _less(self, i, j):
return self.items[i] < self.items[j]
def _swim(self, index):
while index > 0:
parent_index = (index - 1) // 2
if self._less(index, parent_index):
self._exch(index, parent_index)
index = parent_index
else:
break
def _sink(self, index):
while index * 2 + 1 <= len(self.items) - 1:
j = index * 2 + 1
if j < len(self.items) - 1 and self._less(j + 1, j):
j = j + 1
if not self._less(j, index):
break
self._exch(index, j)
index = j
def push(self, item):
self.items.append(item)
self._swim(len(self.items) - 1)
def pop(self):
if not self.items:
raise IndexError('heap is out of range!')
item = self.items[0]
self._exch(0, len(self.items) - 1)
self.items.pop()
self._sink(0)
return item
在二叉堆中添加和删除元素时,为了维持二叉堆的数据结构(确保所以父节点的值比其子节点的值小),当添加元素到堆尾时,若该元素的值小于其父节点元素,则会进行_swim操作。当删除最小元素时,若其值大于其子节点的元素,则会进行_sink操作。维持堆结构的算法时间复杂度为二叉树的高,即为log(N)(N为二叉堆元素个数)。