Maintain a fixed size heap -python

I found this post I was trying to implement a top-n heap with fixed size, here is the solution I can offer:

from heapq import heapify, heappush, heappushpop, nlargest

class MaxHeap():
    def __init__(self, top_n):
        self.h = []
        self.length = top_n
        heapify( self.h)
        
    def add(self, element):
        if len(self.h) < self.length:
            heappush(self.h, element)
        else:
            heappushpop(self.h, element)
            
    def getTop(self):
        return sorted(self.h, reverse=True)

and more optimally (thanks to @CyanoKobalamyne)

from heapq import heapify, heappush, heappushpop, nlargest

class MaxHeap():
    def __init__(self, top_n):
        self.h = []
        self.length = top_n
        heapify( self.h)
        
    def add(self, element):
        if len(self.h) < self.length:
            heappush(self.h, element)
        else:
            heappushpop(self.h, element)
            
    def getTop(self):
        return nlargest(self.length, self.h)

There's no built-in in heapq to check the size, so you'll have to do that yourself:

if len(h) < capacity:
    heapq.heappush(h, thing)
else:
    # Equivalent to a push, then a pop, but faster
    spilled_value = heapq.heappushpop(h, thing)
    do_whatever_with(spilled_value)

Also, note that heapq implements a min heap, not a max heap. You'll need to reverse the order of your priorities, probably by negating them.

Tags:

Python