clrs_heap.py

class Heap:
    """
    Fields:
    -  list _heap
    """
    

    def __init__(self):
        self._heap = []

    @staticmethod
    def _parent(i : int):
        return (i-1) // 2

    @staticmethod
    def _left(i : int):
        return (i << 1) + 1

    @staticmethod
    def _right(i : int):
        return (i << 1) + 2

    def _swap(self, i: int, j: int):
        x = self._heap[i]
        self._heap[i] = self._heap[j]
        self._heap[j] = x
        
    def push(self, x: int):
        self._heap.append(x)
        i = len(self._heap)-1
        while i > 0 and self._heap[self._parent(i)] < self._heap[i]:
            self._swap(self._parent(i),i)
            i = self._parent(i)

    def peek(self):
        """
        Raises [IndexError] exception if the heap is empty.
        """
        return self._heap[0]

    def _max_heapify(self, i : int):
        """
        Assumption: i is the root of a subtree of the heap.
        The left and right subtrees of i are assumed to be heaps already,
        but the heap condition may be violated at node i.
        """

        ll = self._left(i)
        if ll >= len(self._heap):
            return
        largest = i if self._heap[i] >= self._heap[ll] else ll
        rr = self._right(i)
        if rr < len(self._heap) and self._heap[largest] < self._heap[rr]:
            largest = rr
        if largest != i:
            self._swap(i,largest)
            self._max_heapify(largest)
        
    def pop(self):
        """
        Raises [Indexerror] exception if the heap is empty.
        """

        retval = self._heap[0]
        self._heap[0] = self._heap.pop()
        if len(self._heap) > 0:
            self._max_heapify(0)
        return retval

    def validate_min_heap_condition(self):
        for i in range(1,len(self._heap)):
            assert self._heap[self._parent(i)] >= self._heap[i]