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

Python标准库包括heapq module,完成了heapq.nsmallest() function就绪实现：

import heapq

k_smallest = heapq.nsmallest(k,input_list)

在内部,这会创建一个大小为K的堆,其中包含输入列表的前K个元素,然后迭代剩余的N-K个元素,将每个元素推送到堆中,然后弹出最大的元素.这样的推送和弹出需要记录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)