What is the fastest method to get the k smallest numbers in an unsorted list of size N using python?
Is it faster to sort the big list of numbers, and then get the k smallest numbers,
or to get the k smallest numbers by finding the minimum in the list k times, making sure u remove the found minimum from the search before the next search?
解决方案
You could use a heap queue; it can give you the K largest or smallest numbers out of a list of size N in O(NlogK) time.
The Python standard library includes the heapq module, complete with a heapq.nsmallest() function ready implemented:
import heapq
k_smallest = heapq.nsmallest(k, input_list)
Internally, this creates a heap of size K with the first K elements of the input list, then iterating over the remaining N-K elements, pushing each to the heap, then popping off the largest one. Such a push and pop takes log K time, making the overall operation O(NlogK).
The function also optimises the following edge cases:
If K is 1, the min() function is used instead, giving you a O(N) result.
If K >= N, the function uses sorting instead, since O(NlogN) would beat O(NlogK) in that case.
A better option is to use the introselect algorithm, which offers an O(n) option. The only implementation I am aware of is using the 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()
Apart from requiring installation of numpy, this also takes N memory (versus K for heapq), as a copy of the list is created for the partition.
If you only wanted indices, you can use, for either variant:
heapq.nsmallest(k, range(len(input_list)), key=input_list.__getitem__) # O(NlogK)
numpy.argpartition(numpy.array(input_list), k)[:k].tolist() # O(N)