heapq.nlargest的时间复杂度是多少？

Question

我在看这个PyCon演讲，34:30的时候，演讲者提到，从一个包含n个元素的列表中获取t个最大的元素可以在O(t + n)的时间内完成。

这怎么可能呢？我理解的是，创建一个堆的时间复杂度是O(n)，那么nlargest本身的复杂度是O(n + t)还是O(t)呢？（具体的算法是什么？）

Answer 1

根据在heapq源代码中的评论：

# Algorithm notes for nlargest() and nsmallest()
# ==============================================
#
# Make a single pass over the data while keeping the k most extreme values
# in a heap.  Memory consumption is limited to keeping k values in a list.
#

# Theoretical number of comparisons for k smallest of n random inputs:
#
# Step   Comparisons                  Action
# ----   --------------------------   ---------------------------
#  1     1.66 * k                     heapify the first k-inputs
#  2     n - k                        compare remaining elements to top of heap
#  3     k * (1 + lg2(k)) * ln(n/k)   replace the topmost value on the heap
#  4     k * lg2(k) - (k/2)           final sort of the k most extreme values
#
# Combining and simplifying for a rough estimate gives:
#
#        comparisons = n + k * (log(k, 2) * log(n/k) + log(k, 2) + log(n/k))

Answer 2

其实这个算法的时间复杂度是 O(n+tlog(n))。因为“堆化”这个过程需要 O(n) 的时间，而每次处理最大或最小元素需要 O(log(n)) 的时间。所以如果你要找 t 个最大或最小的元素，就需要 tlog(n) 的时间。综上所述，总的时间复杂度就是 O(n+t*log(n))。

Answer 3

对于求最大的 t 个数或最小的 t 个数，时间复杂度是 O(nlog(t))。

Heapq 会先为前 t 个元素建立一个堆，然后再对剩下的元素进行处理，通过不断地把元素放入堆中和从堆中取出元素（保持堆中始终有 t 个元素）。

1. 为前 t 个元素建立堆的过程时间复杂度是 tlog(t)。
2. 对剩下的元素进行放入和取出的操作，时间复杂度是 (n-t)log(t)。
3. 所以，总的时间复杂度是 nlog(t)。

Answer 4

这个发言者在这里说错了。实际的时间复杂度是 O(n * log(t))。堆化（heapify）只在可迭代对象的前 t 个元素上进行。这部分的复杂度是 O(t)，但如果 t 远小于 n，那就不算什么。接下来，剩下的元素通过 heappushpop 一次一个地添加到这个“小堆”里。每次调用 heappushpop 的时间复杂度是 O(log(t))。堆的长度始终保持在 t。最后，堆会被排序，这个过程的复杂度是 O(t * log(t))，但如果 t 远小于 n，这部分也不算什么。

理论的乐趣 ;-)

其实有一些比较简单的方法可以在预期的 O(n) 时间内找到第 t 大的元素；比如，可以参考 这里。还有一些更复杂的方法可以在最坏情况下也做到 O(n) 时间。然后，在对输入数据进行另一遍处理时，你可以输出所有大于等于第 t 大的元素（如果有重复元素的话，处理起来会比较麻烦）。所以整个过程实际上是可以在 O(n) 时间内完成的。

不过，这些方法也需要 O(n) 的内存。Python 并没有使用这些方法。实际上实现的这个方法的一个优点是，最坏情况下“额外”的内存开销是 O(t)，这在输入是一个生成器并且产生很多值的情况下会非常重要。