heapq模块是python的一个标准库,它实现了一个堆数据结构,堆数据结构是一种二叉树。
heap[k] <= heap[2*k+1] and heap[k] <= heap[2*k+2]
for all k, counting elements from zero.
我们可以这样理解:
堆是完全二叉树或者近似二叉树,它的各个节点的键值都有固定对应的的数字,根节点(即root,最上面起始位置)是0,若父节点为heap[k],则子节点为heap[2*k+1]和heap[2*k+2]。父节点对应的值总是小于或者等于子节点,称为最小堆。对应的,父节点的值总是大于或者等于子节点,称为最大堆。在heapq中,使用的是最小堆。
正因为堆的这种特殊结构,使得通过heapq模块,可以快速获取一个列表的前N个最大(小)值,即Top N。
这里,可能我们不禁要问,python不是内置了sort方法用来排序么?
现在我们假设一种情景,我们在维护一个列表,并且这个列表在变化,不断有新元素加入,而在任何时候我们可能需要获取里面的Top N,因此我们要求列表始终可以处于已排序状态。
这时候sort方法就显得不那么好用了,因为每次新加入一个元素,我们都要调用一次sort。数据量小时还是可以的,当数据量较大时,效率就会降低,并且python内置的sort方法本身在实现上也不是那么的高效,复杂度是O(nlog n)。
python维护了一个堆,使用的存储结构是列表,通过heapq模块来管理、操作这个堆。heapq提供了插入、删除元素的方法,并且保证在插入或删除元素时,所有节点自动调整,保证堆的结构,同时尽量高效,复杂度为O(log n),在大数据时,效率高于sort排序。
heapq模块提供了如下几个函数:
heapq.heappush(heap, item) 把item添加到heap中(heap是一个列表)
heapq.heappop(heap) 把堆顶元素弹出,返回的就是堆顶
heapq.heappushpop(heap, item) 先把item加入到堆中,然后再pop,比heappush()再heappop()要快得多
heapq.heapreplace(heap, item) 先pop,然后再把item加入到堆中,比heappop()再heappush()要快得多
heapq.heapify(list) 将列表进行堆调整
heapq.merge(*iterables) 将多个列表合并,并进行堆调整,返回的是合并后的列表的迭代器
heapq.nlargest(n, iterable, key=None) 返回最大的n个元素(Top-K问题)
heapq.nsmallest(n, iterable, key=None) 返回最小的n个元素(Top-K问题)
获取Top N:
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]
使用key关键字参数,进行更为复杂的数据结构处理:
{'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'])
对一个列表进行排序:
heapq.heapify(nums)
print(nums)
#[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
堆数据结构的实现有兴趣的可以查看具体的源码实现,堆本身并不是一种有序的结构,但可以通过遍历二叉树的方式得到有序的列表。
在关于排序和取Top N值时,到底使用什么方法最快,python3 cookbook给出了非常好的建议:
当要查找的元素个数相对比较小的时候,函数nlargest() 和 nsmallest()。
仅仅想查找唯一的最小或最大(N=1)的元素的话,那么使用min()和max()函数。
如果N的大小和集合大小接近的时候,通常先排序这个集合然后再使用切片操作会更快点 (sorted(items)[:N] 或者是 sorted(items)[-N:])。
Leave a Reply