Deap in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/deap/. The authors are listed in prov/5CWRB5.jsonl.

from collections.abc import Iterable
from typing import Any, Generic, Optional, Protocol, TypeVar


class Ordered(Protocol):
    def __lt__(self, other: Any, /) -> bool: ...


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


class Deap(Generic[T]):
    """A deap: a double-ended heap in one 1-based array with an empty root.

    The positions 2 .. n hold the keys. The subtree of 2 is a min-heap and
    the subtree of 3 is a max-heap. Each min node i has a partner j in the
    max-heap at the same level, or the parent of that place if it is past
    n, and A[i] <= A[j]. insert, extract_min and extract_max take
    O(log n) time; find_min and find_max take O(1) time. Keys need only "<".
    """

    def __init__(self, keys: Iterable[T] = ()) -> None:
        self.a: list[Optional[T]] = [None, None]
        for k in keys:
            self.insert(k)

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

    @property
    def n(self) -> int:
        return len(self.a) - 1

    def _k(self, i: int) -> T:
        v = self.a[i]
        assert v is not None
        return v

    @staticmethod
    def is_max_node(i: int) -> bool:
        k = i.bit_length() - 1
        return (i >> (k - 1)) == 3

    @staticmethod
    def min_partner(j: int) -> int:
        return j - (1 << (j.bit_length() - 2))

    def max_partner(self, i: int) -> int:
        j = i + (1 << (i.bit_length() - 2))
        return j if j <= self.n else j // 2

    def _min_sift_up(self, i: int, x: T) -> None:
        a = self.a
        while i > 3 and x < self._k(i // 2):  # @line L6
            a[i] = a[i // 2]
            i //= 2
        a[i] = x  # @line L7

    def _max_sift_up(self, j: int, x: T) -> None:
        a = self.a
        while j > 3 and self._k(j // 2) < x:  # @line L8
            a[j] = a[j // 2]
            j //= 2
        a[j] = x  # @line L9

    def insert(self, x: T) -> None:
        a = self.a
        a.append(x)  # @line L1
        n = self.n
        if n == 2:
            return
        if self.is_max_node(n):  # @line L2
            i = self.min_partner(n)
            if x < self._k(i):  # @line L3
                a[n] = a[i]
                self._min_sift_up(i, x)
            else:
                self._max_sift_up(n, x)
        else:
            j = self.max_partner(n)  # @line L4
            if self._k(j) < x:  # @line L5
                a[n] = a[j]
                self._max_sift_up(j, x)
            else:
                self._min_sift_up(n, x)

    def find_min(self) -> T:
        if self.n < 2:
            raise IndexError("find_min from an empty deap")
        return self._k(2)

    def find_max(self) -> T:
        if self.n < 2:
            raise IndexError("find_max from an empty deap")
        return self._k(3) if self.n >= 3 else self._k(2)

    def extract_min(self) -> T:
        a = self.a
        if self.n < 2:  # @line L10
            raise IndexError("extract_min from an empty deap")
        r = self._k(2)
        x = self._k(self.n)
        a.pop()
        n = self.n
        if n < 2:  # @line L11
            return r
        if n == 2:
            a[2] = x
            return r
        i = 2
        while 2 * i <= n:  # @line L12
            c = 2 * i
            if c + 1 <= n and self._k(c + 1) < self._k(c):
                c += 1
            a[i] = a[c]
            i = c
        j = self.max_partner(i)  # @line L13
        if self._k(j) < x:  # @line L14
            a[i] = a[j]
            self._max_sift_up(j, x)
        else:
            self._min_sift_up(i, x)
        return r

    def extract_max(self) -> T:
        a = self.a
        if self.n < 2:  # @line L15
            raise IndexError("extract_max from an empty deap")
        if self.n == 2:
            v = a.pop()
            assert v is not None
            return v
        r = self._k(3)  # @line L16
        x = self._k(self.n)
        a.pop()
        n = self.n
        if n == 2:  # @line L17
            return r
        j = 3
        while 2 * j <= n:  # @line L18
            c = 2 * j
            if c + 1 <= n and self._k(c) < self._k(c + 1):
                c += 1
            a[j] = a[c]
            j = c
        i = self.min_partner(j)  # @line L19
        if 2 * i <= n:
            c = 2 * i
            if c + 1 <= n and self._k(c) < self._k(c + 1):
                c += 1
            i = c
        if x < self._k(i):  # @line L20
            a[j] = a[i]
            self._min_sift_up(i, x)
        else:
            self._max_sift_up(j, x)
        return r

    def keys(self) -> list[T]:
        return [self._k(i) for i in range(2, self.n + 1)]

    def check_invariant(self) -> None:
        n = self.n
        for i in range(4, n + 1):
            p = i // 2
            if self.is_max_node(i):
                assert not self._k(p) < self._k(i)
            else:
                assert not self._k(i) < self._k(p)
        for i in range(2, n + 1):
            if not self.is_max_node(i):
                j = self.max_partner(i)
                if j >= 3:
                    assert not self._k(j) < self._k(i)