Ladder decomposition in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/ladder-decomposition/. The authors are listed in prov/VJP1A9.jsonl.

from collections import deque


class LadderLevelAncestor:
    """Level ancestor queries by jump pointers and ladders.

    The tree has the vertices 0 .. n-1 (n >= 1), is unweighted, and is given
    by its n - 1 undirected edges and a root. Preprocessing takes
    O(n log n) time and memory (the ladders alone take O(n)); each query
    takes O(1) time. No recursion is used.
    """

    def __init__(self, n: int, edges: list[tuple[int, int]], root: int = 0) -> None:
        if n < 1 or len(edges) != n - 1 or not 0 <= root < n:
            raise ValueError("a tree needs n >= 1 vertices, n - 1 edges and a root in range")
        adj: list[list[int]] = [[] for _ in range(n)]
        for a, b in edges:
            if not (0 <= a < n and 0 <= b < n) or a == b:
                raise ValueError("edge endpoints must be distinct vertices in range")
            adj[a].append(b)
            adj[b].append(a)
        parent = [-1] * n
        depth = [-1] * n
        depth[root] = 0
        order = [root]
        queue = deque([root])
        while queue:  # @line L1
            v = queue.popleft()
            for w in adj[v]:
                if depth[w] == -1:
                    depth[w] = depth[v] + 1
                    parent[w] = v
                    order.append(w)
                    queue.append(w)
        if len(order) != n:
            raise ValueError("the edges do not form a tree")
        height = [1] * n  # vertices on the longest downward path
        long_child = [-1] * n
        for v in reversed(order):  # @line L2
            p = parent[v]
            if p != -1:
                c = long_child[p]
                if c == -1 or height[v] > height[c] or (height[v] == height[c] and v < c):
                    long_child[p] = v
                    height[p] = height[v] + 1
        ladder_of = [0] * n
        top = [0] * n
        ladders: list[list[int]] = []
        for v in order:  # @line L3
            if v != root and long_child[parent[v]] == v:
                continue
            path = []  # @line L4
            x = v
            while x != -1:
                path.append(x)
                x = long_child[x]
            k = len(path)
            above = []  # @line L5
            x = parent[v]
            while x != -1 and len(above) < k:
                above.append(x)
                x = parent[x]
            above.reverse()
            ladder_id = len(ladders)
            ladders.append(above + path)
            for x in path:  # @line L6
                ladder_of[x] = ladder_id
                top[x] = depth[v] - len(above)
        up = [[p if p != -1 else root for p in parent]]
        levels = max(1, (n - 1).bit_length())
        for _ in range(1, levels):  # @line L7
            prev = up[-1]
            up.append([prev[prev[v]] for v in range(n)])
        self.n = n
        self.root = root
        self.parent = parent
        self.depth = depth
        self.height = height
        self.long_child = long_child
        self.ladders = ladders
        self.ladder_of = ladder_of
        self.top = top
        self._up = up

    def level_ancestor(self, v: int, d: int) -> int | None:
        """Return the ancestor of v at depth d, or None if 0 <= d <= depth(v) fails."""
        if not 0 <= v < self.n:
            raise IndexError("vertex out of range")
        if d < 0 or d > self.depth[v]:  # @line L8
            return None
        k = self.depth[v] - d  # @line L9
        if k == 0:
            return v
        i = k.bit_length() - 1  # @line L10
        u = self._up[i][v]
        return self.ladders[self.ladder_of[u]][d - self.top[u]]  # @line L11

    def kth_ancestor(self, v: int, k: int) -> int | None:
        """Return the ancestor of v at distance k, or None if it does not exist."""
        if not 0 <= v < self.n:
            raise IndexError("vertex out of range")
        return self.level_ancestor(v, self.depth[v] - k)