query an immutable tree

query_tree.py

"""
Problem:
Given an immutable tree (where each node is identified by a unique ID and contains a list of children),
write a function to process queries that ask for the n-th parent (iteratively climbing upwards) of a node.
Preprocessing is allowed (in below code, all lines that are used to populating cache are considered as preprocessing),
want fastest average query time, allow some space complexity.

Time complexity: O(log(n))
Space complexity: O(n*log(n))
"""

from typing import Dict, List
from dataclasses import dataclass


@dataclass
class Node:
    id: int
    children: List["Node"]


def query(root: Node, start: int, n: int) -> int | None:
    if n < 0:
        return None

    # node -> [2^0 parent, 2^1 parent, 2^2 parent, ...]
    cache: Dict[int, List[int]] = {}

    # immediate parents
    def dfs(node: Node, parent: Node | None) -> None:
        cache[node.id] = []
        if parent:
            cache[node.id].append(parent.id)
        for child in node.children:
            dfs(child, node)
    dfs(root, None)

    # 2^i parents
    i = 0
    while True:
        done = True

        for node in cache:
            if i < len(cache[node]):
                prev = cache[node][i]
                if i < len(cache[prev]):
                    cache[node].append(cache[prev][i])
                    done = False

        if done:
            break
        i += 1

    if start not in cache:
        return None

    curr = start
    powr = 0
    while n > 0:
        if n & 1:
            if powr >= len(cache[curr]):
                return None
            curr = cache[curr][powr]
        n >>= 1
        powr += 1

    return curr


def main() -> int:
    print(query(
        Node(0, [Node(1, [Node(2, [Node(3, [Node(4, [Node(5, [Node(6, [Node(7, [Node(8, [Node(9, [])])])])])])])])])]),
        9,
        6,
    ))

    return 0


if __name__ == "__main__":
    raise SystemExit(main())