Soft heap in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/soft-heap/. The authors are listed in prov/C1G5PQ.jsonl.

"""Soft heap in the style of Kaplan, Tarjan and Zwick (2013).

Each node of a binary tree holds a list of items, which all carry the same
current key, the ckey of the node. Items in nodes of rank above T may be
corrupted: their ckey is larger than their own key.
"""

import math
from typing import Any, List, Optional, Protocol, Tuple, TypeVar, Generic


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


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


class Node(Generic[K]):
    """A node: ckey, rank, item list and two children."""

    __slots__ = ("ckey", "rank", "items", "left", "right")

    def __init__(self, key: K) -> None:
        self.ckey: K = key
        self.rank = 0
        self.items: List[K] = [key]
        self.left: Optional["Node[K]"] = None
        self.right: Optional["Node[K]"] = None

    def is_leaf(self) -> bool:
        return self.left is None and self.right is None


class SoftHeap(Generic[K]):
    """Soft heap with error rate eps: at most eps * (number of inserts) corrupted items.

    The threshold is T = ceil(lg(1/eps)) + 5, which is what the proof on the
    page needs. A value of t given directly overrides it.
    """

    def __init__(self, eps: float = 0.25, t: Optional[int] = None) -> None:
        if t is None:
            if not 0 < eps < 1:
                raise ValueError("eps must be in (0, 1)")
            t = math.ceil(math.log2(1 / eps)) + 5
        self.t = t
        self.roots: List[Node[K]] = []  # decreasing rank; the end has the smallest rank
        self.best: List[int] = []  # best[i]: index of the root of smallest ckey in roots[0..i]
        self.inserts = 0
        self.count = 0
        self.links = 0
        self.moves = 0

    def __len__(self) -> int:
        return self.count

    def size(self, rank: int) -> int:
        """Target size s_k: 1 for k <= T, then ceil(3 s_{k-1} / 2)."""
        s = 1
        for _ in range(self.t, rank):
            s = (3 * s + 1) // 2
        return s

    def _sift(self, x: Node[K]) -> None:
        """Fill the item list of x from below while it is shorter than s_rank."""
        target = self.size(x.rank)
        while len(x.items) < target and not x.is_leaf():  # @line L3
            if x.left is None or (x.right is not None and x.right.ckey < x.left.ckey):  # @line L4
                x.left, x.right = x.right, x.left
            c = x.left
            assert c is not None
            self.moves += 1
            x.items.extend(c.items)  # @line L5
            x.ckey = c.ckey
            c.items = []
            if c.is_leaf():  # @line L6
                x.left = None
            else:
                self._sift(c)

    def _link(self, x: Node[K], y: Node[K]) -> Node[K]:
        """Make a new node of rank r + 1 over two roots of rank r, and fill it."""
        self.links += 1
        z: Node[K] = Node(x.ckey)  # @line L1
        z.items = []
        z.rank = x.rank + 1
        z.left, z.right = x, y
        self._sift(z)  # @line L2
        return z

    def _fix_best(self, start: int) -> None:
        """Recompute best[start:]."""
        del self.best[start:]
        for i in range(start, len(self.roots)):
            if i == 0 or self.roots[i].ckey < self.roots[self.best[i - 1]].ckey:
                self.best.append(i)
            else:
                self.best.append(self.best[i - 1])

    def _add_root(self, x: Node[K]) -> None:
        """Add a root, linking it with a root of equal rank while there is one."""
        while True:
            i = len(self.roots)
            while i > 0 and self.roots[i - 1].rank < x.rank:
                i -= 1
            if i > 0 and self.roots[i - 1].rank == x.rank:
                x = self._link(x, self.roots.pop(i - 1))
                continue
            self.roots.insert(i, x)
            self._fix_best(min(i, len(self.best)))
            return

    def insert(self, key: K) -> None:
        """Insert key."""
        self.inserts += 1
        self.count += 1
        self._add_root(Node(key))  # @line L7

    def meld(self, other: "SoftHeap[K]") -> None:
        """Move all items of other into this heap. Both must use the same T."""
        if other.t != self.t:
            raise ValueError("the heaps have different thresholds")
        for x in reversed(other.roots):  # @line L8
            self._add_root(x)
        self.inserts += other.inserts
        self.count += other.count
        other.roots, other.best, other.count, other.inserts = [], [], 0, 0

    def find_min(self) -> Tuple[K, K]:
        """Return (item, current key) of an item of smallest current key."""
        if not self.roots:
            raise IndexError("the heap is empty")
        x = self.roots[self.best[-1]]  # @line L9
        return x.items[-1], x.ckey

    def delete_min(self) -> Tuple[K, K]:
        """Remove and return (item, current key) of an item of smallest current key."""
        if not self.roots:
            raise IndexError("the heap is empty")
        i = self.best[-1]
        x = self.roots[i]  # @line L10
        e = x.items.pop()  # @line L11
        ckey = x.ckey
        self.count -= 1
        if 2 * len(x.items) <= self.size(x.rank):  # @line L12
            if not x.is_leaf():
                self._sift(x)  # @line L13
            elif not x.items:
                self.roots.pop(i)  # @line L14
        self._fix_best(min(i, len(self.best)))  # @line L15
        return e, ckey

    # -- checking ----------------------------------------------------------

    def nodes(self) -> List[Node[K]]:
        out: List[Node[K]] = []
        stack = list(self.roots)
        while stack:
            v = stack.pop()
            out.append(v)
            stack.extend(c for c in (v.left, v.right) if c is not None)
        return out

    def corrupted(self) -> int:
        """The number of items whose current key is larger than their own key."""
        return sum(1 for v in self.nodes() for e in v.items if e < v.ckey)

    def check_invariant(self) -> None:
        """Check ranks, heap order of ckeys, item keys, the size bounds and the error bound."""
        ranks = [r.rank for r in self.roots]
        assert all(a > b for a, b in zip(ranks, ranks[1:]))
        assert len(self.best) == len(self.roots)
        for i in range(len(self.roots)):
            m = min(range(i + 1), key=lambda j: (self.roots[j].ckey, j))
            assert not self.roots[m].ckey < self.roots[self.best[i]].ckey
        total = 0
        for v in self.nodes():
            total += len(v.items)
            assert v.is_leaf() or v.items, "an internal node has an empty list"
            assert v.items or v.is_leaf()
            for e in v.items:
                assert not v.ckey < e
            if v.rank <= self.t:
                assert len(v.items) <= 1
            for c in (v.left, v.right):
                if c is not None:
                    assert c.rank == v.rank - 1
                    assert not c.ckey < v.ckey
        assert total == self.count
        assert all(r.items or not r.is_leaf() for r in self.roots)
        assert self.corrupted() * 2**self.t <= 18 * self.inserts