Splay tree in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/splay-tree/. The authors are listed in prov/MX6QFK.jsonl.

"""Splay tree (Sleator and Tarjan 1985) with bottom-up splaying."""

from typing import Any, Iterator


class Node:
    __slots__ = ("key", "value", "left", "right", "p")

    def __init__(self, key: Any, value: Any = None) -> None:
        self.key = key
        self.value = value
        self.left: "Node | None" = None
        self.right: "Node | None" = None
        self.p: "Node | None" = None


def _rotate(x: Node) -> None:
    y = x.p
    assert y is not None
    z = y.p
    if x is y.left:
        y.left = x.right
        if x.right is not None:
            x.right.p = y
        x.right = y
    else:
        y.right = x.left
        if x.left is not None:
            x.left.p = y
        x.left = y
    y.p = x
    x.p = z
    if z is not None:
        if z.left is y:
            z.left = x
        else:
            z.right = x


def _splay(x: Node) -> int:
    """Splay x to the root of its tree; return the number of rotations."""
    rot = 0
    while x.p is not None:
        y = x.p
        z = y.p
        if z is None:
            _rotate(x)
            rot += 1
        elif (x is y.left) == (y is z.left):
            _rotate(y)
            _rotate(x)
            rot += 2
        else:
            _rotate(x)
            _rotate(x)
            rot += 2
    return rot


class SplayTree:
    """Ordered map. Every operation splays, so the amortized cost is O(log n)."""

    def __init__(self) -> None:
        self.root: Node | None = None
        self.n = 0
        self.rotations = 0  # total rotations, for experiments

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

    def _splay(self, x: Node) -> None:
        self.rotations += _splay(x)
        self.root = x

    def access(self, key: Any) -> Node | None:
        """Splay the node of key, or the last node on its search path. Return the node of key or None."""
        x = self.root  # @line L1
        last = None
        while x is not None and (key < x.key or x.key < key):  # @line L2
            last = x
            x = x.left if key < x.key else x.right
        if x is None:  # @line L3
            x = last
        if x is not None:  # @line L4
            self._splay(x)
        r = self.root  # @line L5
        if r is not None and not (key < r.key or r.key < key):
            return r
        return None

    def search(self, key: Any) -> Any:
        x = self.access(key)
        return None if x is None else x.value

    def __contains__(self, key: Any) -> bool:
        return self.access(key) is not None

    def insert(self, key: Any, value: Any = None) -> bool:
        """Map key to value; return True if key was new."""
        if self.root is None:  # @line L6
            self.root = Node(key, value)
            self.n = 1
            return True
        found = self.access(key)  # @line L7
        if found is not None:
            found.value = value
            return False
        z = Node(key, value)  # @line L8
        r = self.root
        assert r is not None
        if key < r.key:  # @line L9
            z.left = r.left
            z.right = r
            r.left = None
        else:  # @line L10
            z.right = r.right
            z.left = r
            r.right = None
        for c in (z.left, z.right):  # @line L11
            if c is not None:
                c.p = z
        self.root = z
        self.n += 1
        return True

    def _join(self, left: Node | None, right: Node | None) -> Node | None:
        """Join two trees whose keys are all smaller in left. Return the new root."""
        if left is None:  # @line L15
            if right is not None:
                right.p = None
            return right
        left.p = None
        m = left
        while m.right is not None:  # @line L16
            m = m.right
        self.rotations += _splay(m)
        m.right = right  # @line L17
        if right is not None:
            right.p = m
        return m

    def delete(self, key: Any) -> bool:
        """Remove key; return True if it was present."""
        x = self.access(key)  # @line L12
        if x is None:
            return False
        left, right = x.left, x.right  # @line L13
        if left is not None:
            left.p = None
        if right is not None:
            right.p = None
        self.root = self._join(left, right)  # @line L14
        self.n -= 1
        return True

    def split(self, key: Any) -> "SplayTree":
        """Keep the keys <= key in self; return a new tree with the keys > key."""
        other = SplayTree()
        self.access(key)  # @line L18
        r = self.root
        if r is None:
            return other
        if not (key < r.key):  # @line L19
            rest = r.right
            r.right = None
        else:  # @line L20
            rest = r
            self.root = r.left
            r.left = None
            if self.root is not None:
                self.root.p = None
        if rest is not None:
            rest.p = None
        other.root = rest
        other.n = sum(1 for _ in other.keys())
        self.n -= other.n
        return other

    def join(self, other: "SplayTree") -> None:
        """Append all keys of other, which must all exceed the keys of self."""
        self.root = self._join(self.root, other.root)
        self.n += other.n
        other.root = None
        other.n = 0

    def minimum(self) -> Any:
        x = self.root
        if x is None:
            return None
        while x.left is not None:
            x = x.left
        self._splay(x)
        return x.key

    def floor(self, key: Any) -> Any:
        """The largest key <= key, or None. Splays."""
        self.access(key)
        r = self.root
        if r is None:
            return None
        if not (key < r.key):
            return r.key
        x = r.left
        if x is None:
            return None
        while x.right is not None:
            x = x.right
        self._splay(x)
        return x.key

    def _nodes(self) -> Iterator[Node]:
        stack: list[Node] = []
        x = self.root
        while stack or x is not None:
            while x is not None:
                stack.append(x)
                x = x.left
            x = stack.pop()
            yield x
            x = x.right

    def keys(self) -> Iterator[Any]:
        for x in self._nodes():
            yield x.key

    def items(self) -> Iterator[tuple[Any, Any]]:
        for x in self._nodes():
            yield x.key, x.value

    def height(self) -> int:
        if self.root is None:
            return -1
        best = 0
        stack = [(self.root, 0)]
        while stack:
            x, d = stack.pop()
            best = max(best, d)
            for c in (x.left, x.right):
                if c is not None:
                    stack.append((c, d + 1))
        return best

    def check_invariant(self) -> None:
        """BST order, parent links, and the size count."""
        if self.root is None:
            assert self.n == 0
            return
        assert self.root.p is None
        ks = list(self.keys())
        assert len(ks) == self.n
        assert all(a < b for a, b in zip(ks, ks[1:]))
        stack = [self.root]
        while stack:
            x = stack.pop()
            for c in (x.left, x.right):
                if c is not None:
                    assert c.p is x
                    stack.append(c)