python - 在 python 中获取大小为 N 的未排序列表中 k 最小数字的最快方法？

Question

使用 python 获取大小为 N 的未排序列表中的 k 个最小数字的最快方法是什么？
是对大数字列表进行排序，然后得到 k 个最小的数字，
还是通过在列表中找到最小值 k 次，确保你从搜索中删除找到的最小值，来获得 k 个最小的数字，这样更快？在下一次搜索之前？

Answer 1

你可以使用堆队列；它可以在 O(NlogK) 时间内为您提供大小为 N 的列表中的 K 个最大或最小数字。

Python 标准库包括 heapq module , 完成 heapq.nsmallest() function准备实现:

import heapq

k_smallest = heapq.nsmallest(k, input_list)

在内部，这会创建一个大小为 K 的堆，其中包含输入列表的前 K 个元素，然后遍历剩余的 N-K 个元素，将每个元素插入堆中，然后弹出最大的一个。这样的 push 和 pop 需要 log K 时间，使得整个操作为 O(NlogK)。

该函数还优化了以下边缘情况:

• 如果 K 为 1，则使用 min() 函数，得到复杂度为 O(N) 的结果。
• 如果 K >= N，则该函数改用排序，因为在这种情况下 O(NlogN) 会优于 O(NlogK)。

更好的选择是使用 introselect algorithm ，它提供了一个 O(n) 选项。我知道的唯一实现是使用 numpy.partition() function :

import numpy

# assuming you have a python list, you need to convert to a numpy array first
array = numpy.array(input_list)
# partition, slice back to the k smallest elements, convert back to a Python list
k_smallest = numpy.partition(array, k)[:k].tolist()

除了需要安装 numpy 之外，这还需要 N 内存(相对于 heapq 的 K)，因为会为分区创建列表的副本。

如果你只想要索引，你可以使用，对于任一变体:

heapq.nsmallest(k, range(len(input_list)), key=input_list.__getitem__)  # O(NlogK)
numpy.argpartition(numpy.array(input_list), k)[:k].tolist()  # O(N)