Randomized heap in Python

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

import random
from typing import Any, Generic, Iterator, Protocol, TypeVar


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


class Coin(Protocol):
    def random(self) -> float: ...


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


class Node(Generic[K]):
    """A node: a key and two subtrees. There is no balance field."""

    __slots__ = ("key", "left", "right")

    def __init__(self, key: K) -> None:
        self.key = key
        self.left: Node[K] | None = None
        self.right: Node[K] | None = None


def meld(a: "Node[K] | None", b: "Node[K] | None", coin: Coin) -> "Node[K] | None":
    """Randomized meld of two heap-ordered trees; return the new root.

    This is the recursive MELD of the page as a loop: at each node of the
    path, a fair coin chooses the child that the rest is melded into.
    """
    if a is None:  # @line L1
        return b
    if b is None:
        return a
    if b.key < a.key:  # @line L2
        a, b = b, a
    root = a
    while True:
        go_left = coin.random() < 0.5  # @line L3
        sub = a.left if go_left else a.right
        if sub is not None and b.key < sub.key:
            sub, b = b, sub  # the smaller root continues the path
        if go_left:  # @line L4
            a.left = sub if sub is not None else b
        else:
            a.right = sub if sub is not None else b
        if sub is None:  # MELD(NIL, b) = b
            break
        a = sub
    return root  # @line L5


class RandomizedHeap(Generic[K]):
    """A min-heap with randomized meld (Gambin and Malinowski, 1998)."""

    def __init__(self, coin: Coin | None = None) -> None:
        self.coin: Coin = coin if coin is not None else random.Random()
        self.root: Node[K] | None = None
        self.n = 0

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

    def is_empty(self) -> bool:
        """Return True if the heap holds no item."""
        return self.root is None

    def insert(self, key: K) -> None:
        """Add key to the heap in O(log n) expected time."""
        self.root = meld(self.root, Node(key), self.coin)  # @line L6
        self.n += 1

    def find_min(self) -> K:
        """Return a smallest key; the heap must not be empty."""
        if self.root is None:
            raise IndexError("find_min on an empty heap")
        return self.root.key

    def delete_min(self) -> K:
        """Remove and return a smallest key in O(log n) expected time."""
        r = self.root
        if r is None:
            raise IndexError("delete_min on an empty heap")
        self.root = meld(r.left, r.right, self.coin)  # @line L7
        self.n -= 1
        return r.key  # @line L8

    def meld(self, other: "RandomizedHeap[K]") -> None:
        """Move all items of other into this heap; other becomes empty."""
        self.root = meld(self.root, other.root, self.coin)
        self.n += other.n
        other.root, other.n = None, 0

    def keys(self) -> Iterator[K]:
        """Yield all keys in preorder."""
        stack = [self.root] if self.root is not None else []
        while stack:
            v = stack.pop()
            yield v.key
            for c in (v.right, v.left):
                if c is not None:
                    stack.append(c)

    def check_invariant(self) -> None:
        """Check heap order and the count."""
        count = 0
        stack = [self.root] if self.root is not None else []
        while stack:
            v = stack.pop()
            count += 1
            for c in (v.left, v.right):
                if c is not None:
                    assert not (c.key < v.key), "heap order"
                    stack.append(c)
        assert count == self.n, "count"