Answer k-step ancestor queries in a rooted tree

Problem Statement

You are given a rooted tree with n nodes, indexed from 0 to n-1. The tree is provided as an array parent of size n, where parent[i] is the parent of node i. The root of the tree is the node where parent[i] == -1.

Implement a data structure that supports the following query: getKthAncestor(node, k), which returns the k-th ancestor of the given node. If the k-th ancestor does not exist, return -1.

Constraints

• 1 <= n <= 10^5
• 1 <= q <= 10^5 (number of queries)
• 0 <= k <= n

Implementation

class TreeAncestor:
    def __init__(self, n: int, parent: list[int]):
        # Initialize the data structure
        pass

    def getKthAncestor(self, node: int, k: int) -> int:
        # Return the k-th ancestor of node
        pass

Example

Input: n = 7, parent = [-1, 0, 0, 1, 1, 2, 2]

• getKthAncestor(3, 1) -> 1 (3's parent is 1)
• getKthAncestor(5, 2) -> 0 (5 -> 2 -> 0)
• getKthAncestor(6, 3) -> -1 (6 -> 2 -> 0 -> None)

Edge Cases

• k = 0 (The ancestor is the node itself).
• k is larger than the depth of the node (return -1).
• The tree is a line (a degenerate linked list).