查找集合中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为二叉堆元素个数)。

posted @ 2019-07-27 10:12  Jeffrey_Yang  阅读(1022)  评论(0编辑  收藏  举报