A minimal example of heap

heap.py

from __future__ import annotations
import heapq
import random
from copy import deepcopy
from typing import List, Callable


class Heap(object):
    def __init__(self, arr: List[int], func: Callable[[int, int], bool]):
        self.arr = arr
        self.func = func        # the comparing strategy between parent node and leaf node
        self.heapify()

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

    def __repr__(self) -> str:
        return str(self.arr)

    def is_empty(self):
        return len(self.arr) == 0

    def _check_index(self, i: int) -> bool:
        return isinstance(i, int) and 0 <= i < len(self.arr)

    @property
    def first_leaf(self) -> int:
        assert not self.is_empty()
        return len(self.arr) // 2

    @property
    def last_leaf(self) -> int:
        assert not self.is_empty()
        return len(self.arr) - 1

    @property
    def last_parent(self):
        assert not self.is_empty()
        return len(self.arr) // 2 - 1

    @property
    def first_parent(self):
        assert not self.is_empty()
        return 0

    def is_leaf(self, i: int) -> bool:
        assert not self.is_empty() and self._check_index(i)
        return i >= self.first_leaf

    def get_parent(self, i: int) -> int:
        assert (not self.is_empty()) and (i != self.first_parent)
        return (i - 1) // 2

    def get_lchild(self, i: int) -> int:
        assert (not self.is_empty()) and (not self.is_leaf(i))
        lc = 2 * i + 1
        return lc if lc < len(self.arr) else None

    def get_rchild(self, i: int) -> int:
        assert (not self.is_empty()) and (not self.is_leaf(i))
        rc = 2 * i + 2
        return rc if rc < len(self.arr) else None

    def _shift_down(self, p: int) -> None:
        if self.is_leaf(p):
            return
        lc, rc = self.get_lchild(p), self.get_rchild(p)
        c = rc if rc is not None and self.func(self.arr[rc], self.arr[lc]) else lc
        if not self.func(self.arr[p], self.arr[c]):
            self.arr[p], self.arr[c] = self.arr[c], self.arr[p]
            self._shift_down(c)

    def _shift_up(self, c: int) -> None:
        if c == self.first_parent:
            return
        p = self.get_parent(c)
        if not self.func(self.arr[p], self.arr[c]):
            self.arr[p], self.arr[c] = self.arr[c], self.arr[p]
            self._shift_up(p)

    def heapify(self) -> None:
        if len(self.arr) == 0:
            return
        for p in range(self.last_parent, self.first_parent - 1, -1):
            self._shift_down(p)

    def _insert(self, val: int) -> None:
        assert isinstance(val, int)
        self.arr.append(val)
        self._shift_up(self.last_leaf)

    def push(self, vals: int | List[int]) -> None:
        if isinstance(vals, int):
            vals = [vals]
        assert isinstance(vals, list) and all(isinstance(v, int) for v in vals)
        for v in vals:
            self._insert(v)

    def pop(self) -> int:
        assert not self.is_empty()
        if len(self.arr) == 1:
            v = self.arr[0]
            self.arr.clear()
            return v
        else:
            self.arr[0], self.arr[-1] = self.arr[-1], self.arr[0]
            v = self.arr.pop(-1)
            self._shift_down(0)
            return v

    def get_top(self) -> int:
        assert not self.is_empty()
        return self.arr[0]


def max_heap_func(v1: int, v2: int) -> bool:
    return v1 >= v2


def min_heap_func(v1: int, v2: int) -> bool:
    return v1 <= v2


def check_heap(h: Heap) -> bool:
    arr, func = h.arr, h.func
    for i in range(0, len(arr) // 2 - 1):
        if not (func(arr[i], arr[2 * i + 1]) and func(arr[i], arr[2 * i + 2])):
            return False
    return True


def heap_sort(arr: List[int]) -> List[int]:
    h = Heap(arr, min_heap_func)
    return [h.pop() for _ in range(len(h))]

Core operations: _shift_down and _shift_up