Treap split in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/treap-split/. The authors are listed in prov/WYJ3AW.jsonl.

from __future__ import annotations

from typing import Generic, Iterator, Protocol, TypeVar


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


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


class Node(Generic[K]):
    __slots__ = ("key", "pri", "left", "right")

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


def build(pairs: list[tuple[K, float]]) -> Node[K] | None:
    """Build the treap of distinct (key, priority) pairs.

    Sorts by key and keeps the right spine on a stack, as for a Cartesian
    tree. O(n log n) for the sort, O(n) after it.
    """
    spine: list[Node[K]] = []
    for k, p in sorted(pairs, key=lambda kp: kp[0]):
        x = Node(k, p)
        last: Node[K] | None = None
        while spine and spine[-1].pri < p:
            last = spine.pop()
        x.left = last
        if spine:
            spine[-1].right = x
        spine.append(x)
    return spine[0] if spine else None


def keys(t: Node[K] | None) -> Iterator[K]:
    """Yield the keys of t in increasing order. O(size of t)."""
    stack: list[Node[K]] = []
    while stack or t is not None:
        while t is not None:
            stack.append(t)
            t = t.left
        t = stack.pop()
        yield t.key
        t = t.right


def check_treap(t: Node[K] | None) -> int:
    """Raise AssertionError unless t has search order on keys and max-heap
    order on priorities; return its number of nodes."""
    count = 0
    stack: list[tuple[Node[K], K | None, K | None]] = [(t, None, None)] if t else []
    while stack:
        x, lo, hi = stack.pop()
        count += 1
        assert lo is None or lo < x.key
        assert hi is None or x.key < hi
        if x.left is not None:
            assert x.left.pri <= x.pri
            stack.append((x.left, lo, x.key))
        if x.right is not None:
            assert x.right.pri <= x.pri
            stack.append((x.right, x.key, hi))
    return count


def split(t: Node[K] | None, k: K) -> tuple[Node[K] | None, Node[K] | None]:
    """Split the treap t into (L, R): the keys less than k go to L, the
    keys at least k go to R. Both are treaps. The nodes of t are reused.

    Precondition: t is a treap. Cost O(depth of the search path of k),
    which is O(log n) expected for random priorities.
    """
    if t is None:  # @line L1
        return None, None
    if t.key < k:  # @line L2
        a, b = split(t.right, k)  # @line L3
        t.right = a  # @line L4
        return t, b  # @line L5
    # @line L6
    a, b = split(t.left, k)  # @line L7
    t.left = b  # @line L8
    return a, t  # @line L9