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

Question

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

Answer 1

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

Python标准库包含该heapq模块,其中包含一个已实现的heapq.nsmallest()功能:

import heapq

k_smallest = heapq.nsmallest(k, input_list)

在内部,这会创建一个大小为K的堆,其中包含输入列表的前K个元素,然后迭代剩余的NK元素,将每个元素推送到堆中,然后弹出最大的元素.这样的推送和弹出需要记录K时间,使得整个操作O(NlogK).

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

• 如果K为1,min()则使用该函数,给出O(N)结果.
• 如果K> = N,则函数使用排序,因为在这种情况下O(NlogN)将击败O(NlogK).

更好的选择是使用introselect算法,它提供O(n)选项.我所知道的唯一实现是使用该numpy.partition()函数:

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个内存(相对于K heapq),因为为分区创建了列表的副本.

如果您只想要索引,则可以使用以下任一变量:

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

Answer 2

如果不需要对第 k 个最小数字的列表进行排序，则可以使用像introselect这样的选择算法在 O(n) 时间内完成。标准库没有，但 NumPy 有numpy.partition这个功能：

partitioned = numpy.partition(l, k)
# The subarray partitioned[:k] now contains the k smallest elements.