Fast way to get N Min or Max elements from a list in Python

Question

I currently have a long list which is being sorted using a lambda function f. I then choose a random element from the first five elements. Something like:

f = lambda x: some_function_of(x, local_variable)
my_list.sort(key=f)
foo = choice(my_list[:4])

This is a bottleneck in my program, according to the profiler. How can I speed things up? Is there a fast, inbuilt way to retrieve the elements I want (in theory shouldn't need to sort the whole list). Thanks.

Answer 1

Use heapq.nlargest or heapq.nsmallest.

For example:

import heapq

elements = heapq.nsmallest(4, my_list, key=f)
foo = choice(elements)

This will take O(N+KlogN) time (where K is the number of elements returned, and N is the list size), which is faster than O(NlogN) for normal sort when K is small relative to N.