题目:
给你一棵树,树上有 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)
提交通过: