Leetcode 1483 Kth Ancestor of a Tree Node

Problem statement

https://leetcode.com/problems/kth-ancestor-of-a-tree-node/

Solution

Let us dp[i][j] be 2^j parent of node i. Then we can do recursion dp[i][j] = dp[dp[i][j-1]][j-1], that is, i-th node’s 2^j parent is i-th node’s 2^{j-1} parent’s 2^{j-1} parent.

Complexity

For init: it is O(n log n) time and space For getKthAncestor time complexity O(log n), space is O(1).

Code

class TreeAncestor:
    def __init__(self, n, parent):
        m = int(log2(n))
        self.dp = [parent] + [[-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 getKthAncestor(self, node, k):
        while k > 0 and node != -1: 
            i = int(log2(k))
            node = self.dp[i][node]
            k -= (1 << i)
        return node