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 。

思路

想找到第 k 个祖先节点，可以一步一步往上跳：

• 第 1 个祖先节点是 parent[i]；
• 第 2 个祖先节点是 parent[parent[i]]；
• 依次类推

这样的话查询第 k 个祖先节点需要 k 次，时间复杂度为 $O(n)$。要进行 n 次查询的时间复杂度也是 $O(n)$，由于查询次数和节点数量都是 50000 次，总的时间复杂度就是 $O(n^2)$，因此很容易就会超出时间限制，那该怎么进行优化呢？

没啥想法，看看题解吧。

原来这也可以动态规划啊，简单来说就是可以预处理出每个节点的第 $2^i$ 个祖先节点，在查询第 k 个祖先节点时，可以按照 2 的幂次将它分解，再去做跳转。

例如，对于 k=13，13 可以分解成 1+4+8，那么就可以先往上跳 1 步，再分别跳 4 步和 8 步。因为我们已经预处理好了第 $2^i$ 个节点，因此这样跳转的时间复杂度就是 $O(\log (n))$。

那么该怎么预处理呢？由于节点编号最大为 50000，因此最多只需要表示 $16$ 个祖先节点（$2^{16}=65536$）。那么我们可以定义一个 ancestor 表，ancestor[i][j] 表示第 i 个节点的第 $2^j$ 个祖先节点。这样的话，就有：

• 对于 j=0，它的父节点就是 parent[i]：ancestor[i][0]=parent[i]；
• 对于 j=1，它的爷爷节点是父节点的父节点，也就是 parent[parent[i]]，带入的话就是 ancestor[i][1]=ancestor[ancestor[i][0]][0]；
• 那么以此类推，可以得出通用的转换方程：ancestor[i][j]=ancestor[ancestor[i][j-1]][j-1]

这样就实现好初始化和预处理的功能了。

接下来是查询 node 第 k 个祖先，思路也很简单：对 k 按位从最低位到最高位进行判断，如果第 j 位为 1，代表节点 node 的第 $2^j$ 的祖先是 ancestor[node][j]，可以得出转移方程 node = ancestor[node][j]。之后依次遍历即可。

在实现中，因为最高次数也不过 16，我们可以对 k 右移最高 16 次，减小代码的复杂度（因为只要 $2^j>node$，那么 node 右移 j 次就一定为 0，不会进行状态转移）。

最终代码为：

class TreeAncestor:
 
    def __init__(self, n: int, parent: List[int]):
        self.LG_N = 16
        self.dp = [[-1]*self.LG_N for _ in range(n)]
        for i in range(n):
            self.dp[i][0] = parent[i]
        for j in range(1, self.LG_N):
            for i in range(n):
                if self.dp[i][j-1] != -1:
                    self.dp[i][j] = self.dp[self.dp[i][j-1]][j-1]
 
    def getKthAncestor(self, node: int, k: int) -> int:
        for j in range(self.LG_N):
            if (k >> j) & 1:
                node = self.dp[node][j]
                if node == -1:
                    return -1
        return node
 

想得太多

为什么初始化时状态转移的实现是

        for j in range(1, self.LG_N):
            for i in range(n):
                if self.dp[i][j-1] != -1:
                    self.dp[i][j] = self.dp[self.dp[i][j-1]][j-1]

而不是

        for i in range(n):
            for j in range(1, self.LG_N):
                if self.dp[i][j-1] != -1:
                    self.dp[i][j] = self.dp[self.dp[i][j-1]][j-1]

呢？

想了一下，没想通，不如直接试一下吧。发现报错：

["TreeAncestor","getKthAncestor"]
[[6,[-1,2,3,4,5,0]],[1,4]]

Wrong: [null,-1]
True:  [null,5]

试一下构造 dp：

-1, -1, -1
2, -1, -1
3, -1, -1
4, -1, -1
5, -1, -1
0, -1, -1

如果是先遍历 i，那么本质上是先更新行，后更新列，对于这种情况，在更新第二个数字的时候就会出现 dp[1][1]=dp[dp[1][0]][0]，此时 dp[1][0]=2，更新后 dp[1][1]=dp[2][0]=3：

-1, -1, -1
2, 3, -1
3, -1, -1
4, -1, -1
5, -1, -1
0, -1, -1

接着看 dp[1][2]=dp[dp[1][1]][1]，dp[1][1]=3，那么 dp[1][2]=dp[3][1]，因为此时 dp[3][1]=-1，还没有更新，这就导致了问题的出现，因此不能先遍历 n。