有没有一种方法可以在Python中对列表进行排序，直到找到第一个已排序的k个元素？

Question

我有一个普通的无序数字列表。我需要在排序后取前k个元素。问题是，如果列表非常长而k很小，则对整个列表进行排序似乎是浪费。我想出了一种算法解决方案，但需要我编写自己的排序实现。我的问题是：是否有一种方法可以使用Python中已经实现的东西来获得相同的效率？

更新： 只是为了澄清，我知道这会给出我需要的答案：sorted(boring_list)[:n]。但我的担忧是效率问题：我不需要为此对整个列表进行排序。

- Sofia Bravo

---

@squiguy：我认为OP正在寻找最小的 n 个元素，而不仅仅是第一个。（由于它是一个列表，如果是这样的话，boring_list[:n]就足够了。） - DSM

@Sofia：这是一个涉及k次迭代的解决方案还是有更好的方法？你可以使用k次迭代进行冒泡排序来实现。通常，归并排序具有n log n的效率，因此如果kn < n log n的情况下，它应该会更好。但是希望你将排序后的项目存储在列表中，以便在下一次迭代中不会重复出现。 - Nishant

@Rob 这是一个从一个固定地理点到其他地理点的距离列表，其他地理点是随机的，所以，是的，我猜这个列表只是随机整数以随机顺序排列。 - Sofia Bravo

@SofiaBravo 先使用 sorted。 - squiguy

@DSM 是的，我需要前n个最小元素。 - Sofia Bravo

Answer 1

你可以使用heapq模块，特别是它的nlargest或nsmallest函数。

或者只需构建堆并调用heappop()。构建堆应该需要 O(n) 的时间，检索 k 个元素需要 O(k*log(n)) 的时间。

---

这里是一个非常简单而小巧的基准测试：

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

如您所见，heapify 在处理包含1000000个元素的列表时比排序要快得多。

- Bakuriu

---

你没有把 heapq.heappop(s) 的结果存储在任何地方。将其存储在列表中会如何影响时间？ - Rob Watts

@RobWatts 没有显著的变化。从2.41毫秒到2.42毫秒，从264毫秒到265毫秒。 - Bakuriu

Answer 2

使用 heapq.nsmallest 函数。

维护堆的不变性是 O(logk)，其中 k 是堆的大小；您必须执行 n 次推操作，使得总体复杂度为 O(n logk)。与排序并取前 k 个元素相比，后者的总体复杂度为 O(n logn)。当 k 相对于 n 很小时，heapq 方法显然胜出。

当 k 接近 n 时，您应该只对其进行排序并取前 k 个元素 - timsort 真的很好：-)

- roippi

Answer 3

我会为此编写自己的函数。

import sys
def sort_first_k(iterable,k):
    lst = [sys.maxsize]
    max_ = (sys.maxsize,0) # (sys.maxint,0) on python2

    for el in iterable:
        if el < max_[0]:
            lst.append(el)
            if len(lst) > k: lst.pop(max_[1])
                tmp = max(lst)
                max_ = (tmp, lst.index(tmp))
    return sorted(lst)

- Adam Smith

Answer 4

为什么要排序？？？这不是作业要求的。

def nsmallest(some_list,N):
    tmp = some_list[:]
    xiter = (x for x in iter(lambda:min(tmp),'') if not tmp.remove(x))
    return [val for i,val in zip(range(N),xiter)]

这应该是O(k*n)

In [52]: the_list = [random.randint(-100,1000) for _ in range(1000000)]

In [53]: %timeit nsmallest(the_list,3)
10 loops, best of 3: 66.9 ms per loop

- Joran Beasley

Answer 5

如果您使用像中位数的中位数这样的选择算法，那么您可以在O(n)时间内获得前k个元素。然后对这k个元素进行排序只需要O(k log k)的时间。因此，所有这些操作总共只需要O(n + k log k)的时间复杂度。

- petabyte

---

这正是我所想的，但我的问题是是否可以使用已经用Python编写的东西来完成。 - Sofia Bravo

O(n + klog(k)) = O(n)。同时，已知中位数算法具有相当大的常数因子，这就是为什么几乎所有快速排序的实现都只使用随机选择而不依赖于它的原因。 - Bakuriu

快速排序有趣的地方在于它的时间复杂度。但是，如何证明 O(n + klog(k)) = O(n) 呢？例如，如果 k = n，那么它的时间复杂度应该是 O(n + n log n)。 - petabyte