Kth Ancestor of a Tree Node LeetCode solution in Python

1
# Runtime: 2359 ms (Top 39.76%) | Memory: 44.3 MB (Top 93.27%)
2
from math import ceil, log2
3
from typing import List
4

5
NO_PARENT = -1
6

7

8
class TreeAncestor:
9
    def __init__(self, n: int, parent: List[int]):
10
        self.parent = [[NO_PARENT] * n for _ in range(ceil(log2(n + 1)))]
11
        self.__initialize(parent)
12

13
    def __initialize(self, parent: List[int]):
14
        self.parent[0], prev = parent, parent
15

16
        for jump_pow in range(1, len(self.parent)):
17
            cur = self.parent[jump_pow]
18

19
            for i, p in enumerate(prev):
20
                if p != NO_PARENT:
21
                    cur[i] = prev[p]
22

23
            prev = cur
24

25
    def getKthAncestor(self, node: int, k: int) -> int:
26
        jump_pow = self.jump_pow
27

28
        while k > 0 and node != NO_PARENT:
29
            jumps = 1 << jump_pow
30

31
            if k >= jumps:
32
                node = self.parent[jump_pow][node]
33
                k -= jumps
34
            else:
35
                jump_pow -= 1
36

37
        return node
38

39
    @property
40
    def jump_pow(self) -> int:
41
        return len(self.parent) - 1