fastest method of getting k smallest numbers in unsorted list of size N in python?

by

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)

More Related Contents:

  1. Find common substring between two strings
  2. Is the time-complexity of iterative string append actually O(n^2), or O(n)?
  3. Sliding window maximum in O(n) time
  4. What is the time complexity of heapq.nlargest?
  5. Finding the most frequent character in a string
  6. Implement pow(x,n)%d with integers only. (without making use of library functions) [duplicate]
  7. How do I profile a Python script?
  8. Cost of len() function
  9. What’s a good rate limiting algorithm?
  10. Pick N distinct items at random from sequence of unknown length, in only one iteration
  11. Find all combinations of a list of numbers with a given sum
  12. Complexity of *in* operator in Python
  13. Why is Bubble Sort implementation looping forever?
  14. Edit Distance in Python
  15. Longest increasing subsequence
  16. Speed of calculating powers (in python)
  17. heapq with custom compare predicate
  18. General bars and stars
  19. Python Integer Partitioning with given k partitions
  20. What is the complexity of the sorted() function?
  21. Non biased return a list of n random positive numbers (>=0) so that their sum == total_sum
  22. “For” loop first iteration
  23. What is time complexity of a list to set conversion?
  24. NumPy grouping using itertools.groupby performance
  25. A set union find algorithm
  26. What is the time complexity of Python List Reverse?
  27. What is the Big O Complexity of Reversing the Order of Columns in Pandas DataFrame?
  28. Python – Removing overlapping lists
  29. What is this odd sorting algorithm?
  30. counting combinations and permutations efficiently

Leave a Comment