max_heap.py

from heapq import heapify, heappush, heappushpop

class MaxHeap:
    def __init__(self, top_n: int):
        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 get_top(self):
        return sorted(self.h, reverse=True)

if __name__ == "__main__":
    # no more than 5 elements
    n_saved = 5
    heap = MaxHeap(top_n=n_saved)
    checkpoints = list(range(10))
    
    for i in checkpoints:
        heap.add(i)

    assert heap.get_top() == checkpoints[:-n_saved-1:-1]