binary_heap.py

from typing import TypeVar, Any, Protocol, Sequence, Generic


class Comparable(Protocol):
    def __lt__(self, param: Any) -> bool:
        ...

    def __gt__(self, param: Any) -> bool:
        ...


T = TypeVar("T", bound=Comparable)


class Heapable(Protocol[T]):
    def __len__(self) -> int:
        ...

    def __getitem__(self, idx: Any) -> T:
        ...

    def __setitem__(self, idx: Any, val: T) -> None:
        ...

    def append(self, val: T) -> None:
        ...

    def pop(self) -> T:
        ...


def compare(a: T, b: T) -> bool:
    return a > b


def heap_heapify_down(heap: Heapable[T], i: int, n: int) -> None:
    while True:
        left = i * 2 + 1
        right = i * 2 + 2
        smallest = i
        if left < n and compare(heap[smallest], heap[left]):
            smallest = left
        if right < n and compare(heap[smallest], heap[right]):
            smallest = right
        if i == smallest:
            break
        heap[smallest], heap[i] = heap[i], heap[smallest]
        i = smallest


def heap_heapify_up(heap: Heapable[T], i: int) -> None:
    while i > 0 and compare(heap[(i - 1) // 2], heap[i]):
        heap[(i - 1) // 2], heap[i] = heap[i], heap[(i - 1) // 2]
        i = (i - 1) // 2


def heap_push(heap: Heapable[T], val: T) -> None:
    heap.append(val)
    heap_heapify_up(heap, len(heap) - 1)


def heap_pop(heap: Heapable[T]):
    heap[0] = heap[-1]
    heap.pop()
    heap_heapify_down(heap, 0, len(heap))


def heap_make(heap: Heapable[T]) -> None:
    for i in range(len(heap) // 2, -1, -1):
        heap_heapify_down(heap, i, len(heap))


def heap_sort(heap: Heapable[T]) -> None:
    N = len(heap)
    heap_make(heap)
    for i in range(N - 1, 0, -1):
        heap[0], heap[i] = heap[i], heap[0]
        heap_heapify_down(heap, 0, i)
    for i in range(N // 2):
        heap[i], heap[N - i - 1] = heap[N - i - 1], heap[i]


def heap_isheap(heap: Heapable[T], idx: int) -> bool:
    if (idx >= len(heap)):
        return True
    left = 2 * idx + 1
    right = 2 * idx + 2
    if left < len(heap) and compare(heap[idx], heap[left]):
        return False
    if right < len(heap) and compare(heap[idx], heap[right]):
        return False
    return heap_isheap(heap, left) and heap_isheap(heap, right)


if __name__ == '__main__':
    DATA = [5, 3, 1, 2, 4, -2, -40, 100]
    DATA_SORTED = [-40, -2, 1, 2, 3, 4, 5, 100]

    def ConstructByPushing(nums: list[int]) -> list[int]:
        heap: list[int] = []
        for num in nums:
            heap_push(heap, num)
        return heap

    def RemoveOneByOne(heap) -> list[int]:
        removed = []
        while heap:
            removed.append(heap[0])
            heap_pop(heap)
        return removed

    def ApplyHeapSort(nums):
        ans = nums[:]
        heap_sort(ans)
        return ans

    def ApplyMakeHeap(nums):
        ans = nums[:]
        heap_make(ans)
        return ans

    assert heap_isheap(ConstructByPushing(DATA), 0) == True
    assert RemoveOneByOne(ConstructByPushing(DATA)) == DATA_SORTED
    assert ApplyHeapSort(DATA) == DATA_SORTED
    assert heap_isheap(ApplyMakeHeap(DATA), 0) == True
    assert RemoveOneByOne(ApplyMakeHeap(DATA)) == DATA_SORTED
    print("OK")