[Heap Implementation] A general heap implementation that can be used as min-heap or max-heap. #data-structures #python

a_heap.py

class Heap:
    def __init__(self, array, is_min_heap=False):
        self.comparison_func = MIN_HEAP_FUNC if is_min_heap else MAX_HEAP_FUNC
        self.heap = self.build_heap(array)
        self.length = len(self.heap)

    def __len__(self):
        return self.length

    def build_heap(self, array):
        parent_idx = (len(array) - 2) // 2
        for idx in reversed(range(parent_idx + 1)):
            self.sift_down(idx, len(array) - 1, array)
        return array

    def sift_down(self, current_idx, end_idx, heap):
        child_one_idx = current_idx * 2 + 1
        while child_one_idx <= end_idx:
            child_two_idx = child_one_idx + 1 if child_one_idx + 1 <= end_idx else -1
            if child_two_idx != -1:
                if self.comparison_func(heap[child_two_idx], heap[child_one_idx]):
                    to_swap_idx = child_two_idx
                else:
                    to_swap_idx = child_one_idx
            else:
                to_swap_idx = child_one_idx
            if self.comparison_func(heap[to_swap_idx], heap[current_idx]):
                self.swap(current_idx, to_swap_idx, heap)
                current_idx = to_swap_idx
                child_one_idx = current_idx * 2 + 1
            else:
                return

    def sift_up(self, current_idx, heap):
        parent_idx = (current_idx - 1) // 2
        while current_idx > 0:
            if self.comparison_func(heap[current_idx], heap[parent_idx]):
                self.swap(current_idx, parent_idx, heap)
                current_idx = parent_idx
                parent_idx = (current_idx - 1) // 2
            else:
                return

    def peek(self):
        return self.heap[0]

    def remove(self):
        self.swap(0, self.length - 1, self.heap)
        to_remove = self.heap.pop()
        self.length -= 1
        self.sift_down(0, self.length - 1, self.heap)
        return to_remove

    def insert(self, value):
        self.heap.append(value)
        self.length += 1
        self.sift_up(self.length - 1, self.heap)

    def swap(self, i, j, array):
        array[i], array[j] = array[j], array[i]


def MAX_HEAP_FUNC(a, b):
    return a > b


def MIN_HEAP_FUNC(a, b):
    return a < b


def main():
    min_heap = Heap([4, 3, 1, 2, 9, 0], True)