DFUDS in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/dfuds/. The authors are listed in prov/BT52AM.jsonl.
"""DFUDS: the depth-first unary degree sequence of an ordinal tree.

U = ( followed by, for each node in preorder, d opening parentheses and
one closing one, where d is its degree. U is balanced, of 2n symbols
counted with the leading '('. A node is named by the position where its
code starts; the root is at 1. findclose, findopen and fwdsearch are
plain scans here; the range min-max tree makes them O(log n) or O(1).
"""

from __future__ import annotations

from typing import Optional, Sequence


class DFUDS:
    """A static ordinal tree in the DFUDS encoding."""

    def __init__(self, children: Sequence[Sequence[int]], root: int = 0) -> None:
        """Build U by preorder; ``self.order`` is the preorder of the input nodes."""
        u = [1]
        self.order: list[int] = []
        stack = [root]
        while stack:
            v = stack.pop()
            self.order.append(v)
            u.extend([1] * len(children[v]))
            u.append(0)
            stack.extend(reversed(children[v]))
        self.U = u
        self.n = len(self.order)
        self.ones = [0]
        for b in u:
            self.ones.append(self.ones[-1] + b)
        self.zeros = [i for i, b in enumerate(u) if b == 0]

    def text(self) -> str:
        """U as a string of parentheses."""
        return "".join("(" if b else ")" for b in self.U)

    # --- primitives -------------------------------------------------------
    def excess(self, i: int) -> int:
        """Opens minus closes in U[0..i]; excess(-1) = 0."""
        return 2 * self.ones[i + 1] - (i + 1)

    def rank0(self, i: int) -> int:
        """Number of ')' in U[0..i-1]."""
        return i - self.ones[i]

    def select0(self, j: int) -> int:
        """Position of the j-th ')', j >= 1; select0(0) = 0 by convention."""
        return self.zeros[j - 1] if j > 0 else 0

    def fwdsearch(self, i: int, d: int) -> int:
        """Least j > i with excess(j) = excess(i) + d."""
        target = self.excess(i) + d
        j = i + 1
        while self.excess(j) != target:
            j += 1
        return j

    def findclose(self, i: int) -> int:
        """The ')' matching the '(' at i."""
        return self.fwdsearch(i, -1)

    def findopen(self, j: int) -> int:
        """The '(' matching the ')' at j."""
        target = self.excess(j)
        i = j - 1
        while self.excess(i) != target:
            i -= 1
        return i + 1

    # --- operations -------------------------------------------------------
    def root(self) -> int:
        """The root node, at position 1."""
        return 1

    def degree(self, v: int) -> int:
        """Number of children of v."""
        return self.select0(self.rank0(v) + 1) - v  # @line L1

    def child(self, v: int, i: int) -> Optional[int]:
        """The i-th child (1-based) of v, or None."""
        if i < 1 or i > self.degree(v):  # @line L2
            return None
        e = self.select0(self.rank0(v) + 1)  # @line L3
        return self.findclose(e - i) + 1  # @line L4

    def parent(self, v: int) -> Optional[int]:
        """The parent of v, or None for the root."""
        if v == 1:  # @line L5
            return None
        o = self.findopen(v - 1)  # @line L6
        return self.select0(self.rank0(o)) + 1  # @line L7

    def child_rank(self, v: int) -> int:
        """The index (1-based) of v among its siblings; v is not the root."""
        o = self.findopen(v - 1)  # @line L8
        return self.select0(self.rank0(o) + 1) - o  # @line L9

    def subtree_size(self, v: int) -> int:
        """Number of nodes in the subtree of v."""
        end = self.fwdsearch(v - 1, -1)  # @line L10
        return (end - v) // 2 + 1  # @line L11

    def next_sibling(self, v: int) -> Optional[int]:
        """The next sibling of v, or None."""
        if v == 1:
            return None
        p = self.parent(v)
        assert p is not None
        return self.child(p, self.child_rank(v) + 1)

    def is_leaf(self, v: int) -> bool:
        """True if v has no children."""
        return self.U[v] == 0

    def preorder_rank(self, v: int) -> int:
        """Preorder number of v, from 0."""
        return self.rank0(v)

    def node_at(self, k: int) -> int:
        """The node with preorder number k."""
        return 1 if k == 0 else self.select0(k) + 1