LeetCode每日一题(20240406) - 树节点的第K个祖先

import math


class TreeAncestor:

    def __init__(self, n: int, parent: List[int]):
        self._l = math.ceil(math.log2(n))
        # st[i][j]: 第i个节点跳2^j步所能到达的祖先节点
        self._st = [[-1] * self._l for _ in range(n)]
        # st[i][0]即为父节点
        for i in range(n):
            self._st[i][0] = parent[i]
        for j in range(1, self._l):
            for i in range(n):
                prev_parent = self._st[i][j - 1]
                if prev_parent != -1:
                    self._st[i][j] = self._st[prev_parent][j - 1]

    def getKthAncestor(self, node: int, k: int) -> int:
        curr_node = node
        step = 0  # 2^step步
        while k and curr_node != -1:
            if step >= self._l:
                return -1
            if k & 1:  # 比较位来取决于要不要跳
                curr_node = self._st[curr_node][step]
            k >>= 1
            step += 1
        return curr_node