Heap Implementation

Heap.py

class Heap:
    def __init__(self, key, arr=[]) -> None:
        self.arr = arr
        self.key = key

    def add(self, element) -> None:
        self.arr.append(element)
        self.heapify_up()

    def pop(self):
        first = self.arr[0]
        if len(self.arr) > 1:
            self.heapify_down()
        else:
            self.arr.pop()
        return first

    def heapify_up(self) -> None:
        index = len(self.arr) - 1
        while index > 0:
            is_left = int(index % 2 == 1)
            parent_index = (index - is_left) // 2
            if self._compare(index, parent_index):
                self._swap(index, parent_index)
                index = parent_index
            else:
                break

    def heapify_down(self) -> None:
        element = self.arr.pop()
        index = 0
        self.arr[index] = element
        length = len(self.arr)
        while index < length:
            left_child_index = 2 * index + 1
            right_child_index = 2 * index + 2
            exists_left = left_child_index < length
            exists_right = right_child_index < length
            is_left_child = exists_left and self._compare(
                left_child_index, right_child_index)
            if is_left_child and exists_left and self._compare(left_child_index, index):
                self._swap(index, left_child_index)
                index = left_child_index
            elif (not is_left_child) and exists_right and self._compare(right_child_index, index):
                self._swap(index, right_child_index)
                index = right_child_index
            else:
                break

    def _swap(self, index_one: int, index_two: int) -> None:
        self.arr[index_one], self.arr[index_two] = self.arr[index_two], self.arr[index_one]

    def _compare(self, index_one: int, index_two: int) -> bool:
        if index_two >= len(self.arr):
            return True
        return self.key((self.arr[index_one], self.arr[index_two]))

    def __len__(self) -> int:
        return len(self.arr)