A rooted tree has nodes 0..n-1. You get it as a list parent where parent[i] is the parent of node i and parent[root] == -1 (the root is node 0). Labels are arbitrary: a parent's number may be larger than its child's.
Implement a class:
TreeAncestor(n: int, parent: list[int])preprocesses the tree.get_kth_ancestor(node: int, k: int) -> intreturns the ancestorksteps abovenode(k = 1is the parent), or-1if the path to the root is shorter thank.
# 0
# / \
# 2 1
# |
# 3
t = TreeAncestor(4, [-1, 0, 0, 2])
t.get_kth_ancestor(3, 1) # 2
t.get_kth_ancestor(3, 2) # 0
t.get_kth_ancestor(3, 3) # -1
t.get_kth_ancestor(1, 0) # 1 (zero steps is the node itself)
Constraints: 1 <= n <= 5 * 10^4, 0 <= k <= n, up to 5 * 10^4 queries.
Walking up one parent at a time is O(n) per query and far too slow on a long chain. Aim for O(n log n) preprocessing and O(log n) per query.
Show hint
Any k is a sum of powers of two. If you knew, for every node, its ancestor 1, 2, 4, 8, ... steps up, a query would take one jump per set bit of k, and each of those tables can be built from the previous one.