在 Python 中查找树节点的第 K 个祖先的程序

假设我们有一棵树,其中有 n 个节点,从 0 到 n-1 编号。树由一个父数组给出,其中 parent[i] 是节点 i 的父节点。树的根是节点 0。我们必须找到给定节点的第 k 个祖先,如果祖先不存在,则返回 -1

所以,如果输入像

那么输出将是 2,因为节点 6 的第一个祖先是 5,第二个是 2。

示例

让我们看看以下实现以更好地理解

def solve(parent, node, k):
   if node == -1:
      return -1
   elif k == 1:
      return parent[node]
   elif not (k & k-1):
      return solve(parent, solve(parent, node, k >> 1), k >> 1)
   else:
      msb = 1 << (k.bit_length()-1)
      return solve(parent, solve(parent, node, k-msb), msb)

parent = [-1,0,0,1,2,2,5,5]
node = 6
k = 2
print(solve(parent, node, k))

输入

[6,7,9,16,22], 2
输出结果
2