Hybrid ternary search trie with R^2 root branching in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/hybrid-tst/. The authors are listed in prov/EHWQ99.jsonl.

from typing import Generic, TypeVar

V = TypeVar("V")


class _Node(Generic[V]):
    __slots__ = ("c", "left", "mid", "right", "val")

    def __init__(self, c: str) -> None:
        self.c = c
        self.left: _Node[V] | None = None
        self.mid: _Node[V] | None = None
        self.right: _Node[V] | None = None
        self.val: V | None = None


class _Cell(Generic[V]):
    """One entry of the R by R table: the value of the 2-symbol key and the TST of the rest."""

    __slots__ = ("val", "tst")

    def __init__(self) -> None:
        self.val: V | None = None
        self.tst: _Node[V] | None = None


class HybridTST(Generic[V]):
    """A ternary search tree whose root is an R by R table indexed by the first two symbols.

    The alphabet is a string of R distinct symbols in increasing order. Keys of length 0 and 1 are kept
    in a field and an array of R values; a key of length 2 is kept in its
    table cell. get and put take O(1) for the table plus a TST search on the
    remaining m - 2 symbols. Values must not be None.
    """

    def __init__(self, alphabet: str = "".join(map(chr, range(256)))) -> None:
        if list(alphabet) != sorted(set(alphabet)):
            raise ValueError("the alphabet must be sorted and without repeats")
        self.alphabet = alphabet
        self.code = {c: i for i, c in enumerate(alphabet)}
        r = self.r = len(alphabet)
        self.empty_val: V | None = None
        self.one: list[V | None] = [None] * r
        self.table: list[_Cell[V] | None] = [None] * (r * r)
        self.n = 0

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

    def _code(self, c: str) -> int:
        k = self.code.get(c)
        if k is None:
            raise ValueError(f"symbol {c!r} is not in the alphabet")
        return k

    def _cell_index(self, s: str) -> int:
        return self._code(s[0]) * self.r + self._code(s[1])

    def get(self, s: str) -> V | None:
        """Return the value of key s, or None."""
        if len(s) < 2:  # @line L1
            if not s:
                return self.empty_val
            return self.one[self._code(s[0])]
        cell = self.table[self._cell_index(s)]  # @line L2
        if cell is None:  # @line L3
            return None
        if len(s) == 2:  # @line L4
            return cell.val
        x, d, rest = cell.tst, 0, s[2:]  # @line L5
        while x is not None:
            if rest[d] < x.c:
                x = x.left
            elif rest[d] > x.c:
                x = x.right
            elif d < len(rest) - 1:
                x, d = x.mid, d + 1
            else:
                return x.val
        return None

    def put(self, s: str, v: V) -> None:
        """Set the value of key s to v (v must not be None)."""
        if v is None:
            raise ValueError("a value must not be None")
        if len(s) < 2:  # @line L6
            if not s:
                old = self.empty_val
                self.empty_val = v
            else:
                k = self._code(s[0])
                old = self.one[k]
                self.one[k] = v
            if old is None:
                self.n += 1
            return
        i = self._cell_index(s)  # @line L7
        cell = self.table[i]
        if cell is None:
            cell = self.table[i] = _Cell()
        if len(s) == 2:  # @line L8
            if cell.val is None:
                self.n += 1
            cell.val = v
            return
        rest = s[2:]  # @line L9
        if cell.tst is None:
            cell.tst = _Node(rest[0])
        x, d = cell.tst, 0
        while True:
            c = rest[d]
            if c < x.c:
                if x.left is None:
                    x.left = _Node(c)
                x = x.left
            elif c > x.c:
                if x.right is None:
                    x.right = _Node(c)
                x = x.right
            elif d < len(rest) - 1:
                d += 1
                if x.mid is None:
                    x.mid = _Node(rest[d])
                x = x.mid
            else:
                if x.val is None:
                    self.n += 1
                x.val = v
                return

    def delete(self, s: str) -> None:
        """Remove key s if present; prune dead TST nodes and empty cells."""
        if len(s) < 2:  # @line L10
            if not s:
                if self.empty_val is not None:
                    self.empty_val = None
                    self.n -= 1
            else:
                k = self._code(s[0])
                if self.one[k] is not None:
                    self.one[k] = None
                    self.n -= 1
            return
        i = self._cell_index(s)  # @line L11
        cell = self.table[i]
        if cell is None:
            return
        if len(s) == 2:  # @line L12
            if cell.val is not None:
                cell.val = None
                self.n -= 1
        else:
            self._tst_delete(cell, s[2:])  # @line L13
        if cell.val is None and cell.tst is None:  # @line L14
            self.table[i] = None

    def _tst_delete(self, cell: _Cell[V], rest: str) -> None:
        path: list[tuple[_Node[V], _Node[V] | None, str]] = []
        x, parent, link, d = cell.tst, None, "root", 0
        while x is not None:
            path.append((x, parent, link))
            if rest[d] < x.c:
                x, parent, link = x.left, x, "left"
            elif rest[d] > x.c:
                x, parent, link = x.right, x, "right"
            elif d < len(rest) - 1:
                x, parent, link, d = x.mid, x, "mid", d + 1
            else:
                break
        if x is None or x.val is None:
            return
        x.val = None
        self.n -= 1
        for node, par, lk in reversed(path):
            if node.val is not None or node.mid is not None:
                return
            repl = _bst_remove_root(node)
            if par is None:
                cell.tst = repl
            else:
                setattr(par, lk, repl)

    def keys_with_prefix(self, p: str) -> list[str]:
        """Return the keys that start with p, in lexicographic order."""
        out: list[str] = []
        if p == "":  # @line L15
            if self.empty_val is not None:
                out.append("")
            for a in range(self.r):
                out.extend(self.keys_with_prefix(self.alphabet[a]))
            return out
        if len(p) == 1:  # @line L16
            a = self._code(p)
            if self.one[a] is not None:
                out.append(p)
            for b in range(self.r):
                cell = self.table[a * self.r + b]
                if cell is not None:
                    two = p + self.alphabet[b]
                    if cell.val is not None:
                        out.append(two)
                    _collect(cell.tst, two, out)
            return out
        cell = self.table[self._cell_index(p)]  # @line L17
        if cell is None:
            return out
        if len(p) == 2:
            if cell.val is not None:
                out.append(p)
            _collect(cell.tst, p, out)
            return out
        rest = p[2:]  # @line L18
        x, d = cell.tst, 0
        while x is not None:
            if rest[d] < x.c:
                x = x.left
            elif rest[d] > x.c:
                x = x.right
            elif d < len(rest) - 1:
                x, d = x.mid, d + 1
            else:
                break
        if x is None:
            return out
        if x.val is not None:
            out.append(p)
        _collect(x.mid, p, out)
        return out

    def keys(self) -> list[str]:
        return self.keys_with_prefix("")

    def longest_prefix_of(self, s: str) -> str | None:
        """Return the longest key that is a prefix of s, or None."""
        best = 0 if self.empty_val is not None else -1
        if len(s) >= 1 and self.one[self._code(s[0])] is not None:
            best = 1
        if len(s) >= 2:
            cell = self.table[self._cell_index(s)]
            if cell is not None:
                if cell.val is not None:
                    best = 2
                x, d = cell.tst, 2
                while x is not None and d < len(s):
                    if s[d] < x.c:
                        x = x.left
                    elif s[d] > x.c:
                        x = x.right
                    else:
                        d += 1
                        if x.val is not None:
                            best = d
                        x = x.mid
        return None if best < 0 else s[:best]

    def tst_node_count(self) -> int:
        count = 0
        for cell in self.table:
            if cell is not None:
                stack = [cell.tst]
                while stack:
                    y = stack.pop()
                    if y is not None:
                        count += 1
                        stack.extend((y.left, y.mid, y.right))
        return count

    def check_invariant(self) -> None:
        valued = (self.empty_val is not None) + sum(v is not None for v in self.one)
        for cell in self.table:
            if cell is None:
                continue
            assert cell.val is not None or cell.tst is not None, "empty cell"
            valued += cell.val is not None
            stack: list[tuple[_Node[V] | None, str | None, str | None]] = [(cell.tst, None, None)]
            while stack:
                y, lo, hi = stack.pop()
                if y is None:
                    continue
                assert (lo is None or lo < y.c) and (hi is None or y.c < hi)
                assert y.val is not None or y.mid is not None, "dead node"
                valued += y.val is not None
                stack.append((y.left, lo, y.c))
                stack.append((y.right, y.c, hi))
                stack.append((y.mid, None, None))
        assert valued == self.n


def _bst_remove_root(node: _Node[V]) -> _Node[V] | None:
    if node.left is None:
        return node.right
    if node.right is None:
        return node.left
    p, m = node, node.right
    while m.left is not None:
        p, m = m, m.left
    if p is not node:
        p.left = m.right
        m.right = node.right
    m.left = node.left
    return m


def _collect(x: _Node[V] | None, prefix: str, out: list[str]) -> None:
    """Append the keys of the TST below x, each with the given prefix, in order."""
    stack: list[tuple[_Node[V] | None, str, int]] = [(x, prefix, 0)]
    while stack:
        y, t, state = stack.pop()
        if y is None:
            continue
        if state == 0:
            stack.append((y.right, t, 0))
            stack.append((y, t, 1))
            stack.append((y.left, t, 0))
        else:
            if y.val is not None:
                out.append(t + y.c)
            stack.append((y.mid, t + y.c, 0))