初学python记录：力扣1483. 树节点的第 K 个祖先

题目：

给你一棵树，树上有 n 个节点，按从 0 到 n-1 编号。树以父节点数组的形式给出，其中 parent[i] 是节点 i 的父节点。树的根节点是编号为 0 的节点。

树节点的第 k 个祖先节点是从该节点到根节点路径上的第 k 个节点。

实现 TreeAncestor 类：

• TreeAncestor（int n， int[] parent） 对树和父数组中的节点数初始化对象。
• getKthAncestor(int node, int k) 返回节点 node 的第 k 个祖先节点。如果不存在这样的祖先节点，返回 -1 。

思考：

1.暴力解法：

对节点node，在parent数组中向前逐个溯源，找到第k个祖先节点。代码如下：

class TreeAncestor(object):

    def __init__(self, n, parent):
        """
        :type n: int
        :type parent: List[int]
        """
        self.n = n
        self.parent = parent


    def getKthAncestor(self, node, k):
        """
        :type node: int
        :type k: int
        :rtype: int
        """
        for count in range(0, k):
            index = self.parent[node]
            node = index
            if index == -1:
                break
        return index



# Your TreeAncestor object will be instantiated and called as such:
# obj = TreeAncestor(n, parent)
# param_1 = obj.getKthAncestor(node,k)

超时了，卡在8 / 17 个例子。

2. 树上倍增算法

参考. - 力扣（LeetCode）

预处理出每个节点的第 2^i 个祖先节点，即第 1,2,4,8,⋯ 个祖先节点（其中 x 的第 1 个祖先节点就是 parent[x]）。由于任意 k 可以分解为若干不同的 2 的幂（例：13=8+4+1），所以只需要预处理出这些 2^i 祖先节点，就可以快速地到达任意第 k 个祖先节点。

例如 k=13=8+4+1=1101(2) ，我们可以先往上跳 8 步，再往上跳 4 步和 1 步。无论如何跳，都只需要跳 3 次就能到达第 13个祖先节点。

1. 计算祖先节点：

使用数组 pa[x][i] 表示节点x的第 2^i 个祖先节点，i+1≤log2(n)

pa[x][0] = parent[x]     # x的父亲节点

pa[x][1] = pa[pa[x][0]][0]      # x的爷爷节点

pa[x][2] = pa[pa[x][1]][1]

……

pa[x][i+1] = pa[pa[x][i]][i]      # x的第 2^(i+1) 个祖先节点即x的第 2^i 个祖先节点的第 2^i 个祖先节点。另外，若pa[x][i]=-1，则pa[x][i+1]=-1。

2. 拆分k：

从低位（0）到高位判断k的二进制的每一位，若第i位是1，则往上跳 2^i 步，然后判断下一位；若第i位是0，则继续判断下一位。

代码如下（千万注意更新pa数组的代码里，不要写反两个循环的嵌套关系.....血泪的教训T T）：

class TreeAncestor(object):

    def __init__(self, n, parent):
        """
        :type n: int
        :type parent: List[int]
        """
        m = n.bit_length() - 1     # log2(n)=n的二进制位数减一
        pa = [[p] + [-1] * m for p in parent]    # 包含 parent 列表中每个元素及 m 个 -1 组成的列表
        for i in range(0, m):
            for x in range(0, n):     # 循环嵌套关系千万不能变！！否则会出现赋值的时候，pa[pa[x][i]][i]并没有更新的情况，导致pa数组部分值没有更新，而是保持初始值-1
                if pa[x][i] != -1:
                    pa[x][i+1] = pa[pa[x][i]][i]
        self.pa = pa


    def getKthAncestor(self, node, k):
        """
        :type node: int
        :type k: int
        :rtype: int
        """
        n = k.bit_length()    # k的二进制位数
        for i in range(0, n):
            if (k >> i) & 1 == 1:
                node = self.pa[node][i]
                if node < 0:
                    break
        return node


# Your TreeAncestor object will be instantiated and called as such:
# obj = TreeAncestor(n, parent)
# param_1 = obj.getKthAncestor(node,k)

提交通过：