CODEWITHSUNDEEP
pythonalgorithmheapheapsort
Lin Ma
Lin Ma

Reputation: 10139

heap sort in Python for advice

I only need to retrieve the 3 smallest elements, and wondering if there is a way to improve my below code to keep heap size smaller -- I think if we only need to keep heap size as 3, it is enough. But cannot find an option in heapq to tweak.

In other words, I want to want to maintain a three element heap that is occasionally updated.

import heapq

def heapsort(iterable):
   h = []
   for value in iterable:
       heapq.heappush(h, value)
   return [heapq.heappop(h) for i in range(len(h))]

if __name__ == "__main__":

   print heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])

Upvotes: 0

Views: 138

Answers (1)

ShadowRanger
ShadowRanger

Reputation: 155363

The way to improve the code to only get the three smallest elements is to replace it with heapq.nsmallest:

print heapq.nsmallest(3, [1, 3, 5, 7, 9, 2, 4, 6, 8, 0])

Output:

[0, 1, 2]

You can look at the implementation of nsmallest if you're curious about how you'd build it from the heapq primitive functions, because they did exactly that.

Upvotes: 1

Related Questions