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.