从 heapq 到 TopK
最近在看 Python3 Cookbook,还挺有趣的。
Heapq
寻找最小/最大元素
其中 1.4 节提到了一个问题「查找最大或最小的 N 个元素」,用到了内置的 heapq
模块的两个函数 nlargest()
和 nsmallest()
。本质上是利用了最小/最大堆,从堆顶得到堆里面最小/最大的元素。
查看 CPython
源码,nsmallest()
一段典型的实现是这样的:
|
|
这里考虑没有 key
参数的时候的普遍情况。
_heapify_max
通过一个in-place
的方法把包含前 k 个元素和序号的元组列表result
最大堆化。_heapreplace
是最大堆上heappop + heappush
。对所有的元素循环,如果元素比堆顶小,那么堆顶不配在结果里面,把堆顶弹出,然后新元素的元组加进去。- 最后对 k 个元素排序,就可以输出。
另外在代码的注释里面,还给出了各个步骤进行比较的次数。
|
|
建堆
首先第一步建堆,一个 bottom up
的建堆过程可以描述为:
|
|
其中 _siftup
即为《算法导论》中提到的 MIN-HEAPIFY
操作,目的是保持堆的性质,使得以 i
为根的堆满足最小堆的性质。因为大于循环范围的结点都是叶子节点,已经满足堆性质。
这里如果仔细读了 _siftup
源码及注释可以发现,实际实现和书上的方法不太一样:
MIN-HEAPIFY
在大多书上会在当前节点满足小于两个子节点就继续循环_siftup
会把当前子树的值最小的节点上浮
为什么这么做,主要原因是 heappop
也用到了这个内部函数,而当 pop
再调整的时候,最末的节点会移动到堆顶,这个值一般都要比两个子节点大,因此并不能提前终止循环。而对于这个过程中,因为节点的值可能是元组,比较操作的代价比较昂贵,用后面的方法可以有效减少 heappop
比较的次数。
复杂度
堆调整的复杂度可以通过以下方法分析得到:
$$T(k) \le T(2k/3) + \Theta(1)$$其中 $2n/3$ 是最坏情况,及底层恰好半满的时候,递归调用该子树时最大的节点数量。最后可以得到调整的复杂度为 $O(\lg n) = O(h)$。
另外,建堆的复杂度,按照 MIN-HEAPIFY
的方法,在高度 h
上(高度定义为从本节点到叶子节点的最长简单下降路径的边的数量),最多有 $\lceil \frac{k}{2^{h+1}} \rceil$ 个节点,那么可以计算总复杂度为
同时看到源码注释里面给出的是 1.66*k
的比较次数。
整体的复杂度就是 $O(n\lg k)$。
快速选择算法
有没有更快的算法呢?有的。类似于快速排序的快速选择算法。
主要采用了快排的切分的操作,但是只迭代符合要求的一侧元素,以下是一个比较简答的 Python 实现。
|
|
因为每次平均只需要遍历之前数组长度的一半,因此平均复杂度为 $O(n)$,当然最坏的复杂度还是 $O(n^2)$,关键还是看基准值的选择。因此针对此算法的优化变体基本上围绕基准值的选择,比如随机、 BFPRT 。
另外还有一些并行的方法,比如这篇论文 。