Problem statement

https://binarysearch.com/problems/Ancestor-Queries/

Solution

Equal to 1483. Kth Ancestor of a Tree Node, but first create parents array.

Complexity

It is O(n log n) time/space for init an O(log n)/O(1) for query.

Code

class AncestorQuerier:
    def __init__(self, root):
        self.pars, self.vals, self.d = [], {}, []
        def dfs(node, par):
            if not node: return
            self.vals[node.val] = len(self.vals)
            self.d += [node.val]
            self.pars += [self.vals[par.val] if par else -1]
            dfs(node.left, node)
            dfs(node.right, node)

        dfs(root, None)
        n = len(self.pars)
        m = int(log2(n))

        self.dp = [self.pars] + [[-1] * n for _ in range(m)]
        for j in range(m):
            for i in range(n):
                if self.dp[j][i] != -1: 
                    self.dp[j+1][i] = self.dp[j][self.dp[j][i]]

    def getAncestor(self, node, k):
        node = self.vals[node]
        while k > 0 and node != -1: 
            i = int(log2(k))
            node = self.dp[i][node]
            k -= (1 << i)
        return self.d[node] if node != -1 else -1