python - 有没有办法在找到第一个排序的 k 元素之前在 python 中对列表进行排序？

Question

我有一个普通无聊的未排序数字列表。从该列表中，我需要在排序后取出前 k 个元素。问题是，如果列表相当长而 k 相当小，则对整个列表进行排序似乎是一种浪费。我为此想出了一个算法解决方案，但需要我编写自己的排序实现，我的问题是:有没有办法使用已经在 python 中实现的东西获得相同的效率？

更新:

澄清一下，我知道这会给出我需要的答案:sorted(boring_list)[:n]

但我关心的是效率:我不需要为此对整个列表进行排序。

Answer 1

您可以使用 heapq模块，特别是它的 nlargest或 nsmallest功能。

或者只构建堆并调用 heappop() .这应该花费 O(n) 的时间来构建堆和 O(k*log(n)) 来检索 k 元素。

---

这是一个非常简单的小型基准测试:

In [1]: import random, heapq

In [2]: seq = [random.randint(-5000, 5000) for _ in range(35000)]

In [3]: %timeit sorted(seq)[:75]
100 loops, best of 3: 14.5 ms per loop

In [4]: %%timeit
   ...: s = seq[:]
   ...: heapq.nsmallest(75, s)
   ...: 
100 loops, best of 3: 4.05 ms per loop

In [5]: %%timeit
   ...: s = seq[:]
   ...: heapq.heapify(s)
   ...: for _ in range(75): heapq.heappop(s)
   ...: 
100 loops, best of 3: 2.41 ms per loop

我不知道为什么 nsmallest 比直接调用 heappop 慢得多。事实上，我应该在不复制 seq 的情况下计时，但仍然:

In [6]: %%timeit
   ...: heapq.nsmallest(75, seq)
   ...: 
100 loops, best of 3: 3.82 ms per loop

将长度增加 100 倍:

In [12]: %timeit sorted(seq)[:75]
1 loops, best of 3: 1.9 s per loop

In [13]: %%timeit
    ...: heapq.nsmallest(75, seq)
    ...: 
1 loops, best of 3: 352 ms per loop

In [14]: %%timeit
    ...: s = seq[:]
    ...: heapq.heapify(s)
    ...: for _ in range(75): heapq.heappop(s)
    ...: 
1 loops, best of 3: 264 ms per loop

---

注意:为了对抗 F.J 偏见的分析:

In [13]: a = list(range(1000000))

In [14]: random.shuffle(a)

In [15]: %timeit sorted(a)
1 loops, best of 3: 985 ms per loop

In [16]: %%timeit
    ...: s = a[:]
    ...: heapq.heapify(s)
    ...: 
1 loops, best of 3: 284 ms per loop

如您所见，即使在包含 1000000 个元素的列表上，heapify 也比排序快得多。