这道题是最近公共祖先问题(LCA)的拓展。
class FindDistance:
"""
1740. 找到⼆叉树中的距离
https://leetcode.cn/problems/find-distance-in-a-binary-tree/
"""
def solution(self, root, p, q):
self.found = False
self.res = 0
self.lca(root, p, q)
return self.res
# 定义:当⼦树中不包含 p 或 q 时,返回 0;
# 当⼦树中仅包含 p 或 q 中的⼀个时,返回 root 到 p 或 q 的距离;
# 当⼦树中同时包含 p 和 q 时,返回⼀个⽆意义的值(答案会被存在外部变量 res 中)
def lca(self, root, p, q):
if self.found:
# found 为 true 时答案已经被记录在全局 res 中
# 递归函数的返回值已不需要了,此时返回什么值都⽆所谓
return 123
if not root:
return 0
left = self.lca(root.left, p, q)
right = self.lca(root.right, p, q)
# root 的左右⼦树都不包含 p 或 q
if left == 0 and right == 0:
if root.val == p or root.val == q:
return 1
return 0
# 当前节点 root 就是 p, q 的最近公共祖先节点
# 答案已经算出来了,更新全局变量
if left != 0 and right != 0 and not self.found:
self.res = left + right
# 这个递归函数的返回值已经不重要了,让它终⽌递归
self.found = True
return 456
# 此时 left 和 right 有⼀个为 0,即只找到了⼀个节点
# branch 就是到该节点的距离
branch = right if left == 0 else left
# 此时root就是p和q的最近祖先节点了,可直接更新res
if root.val == p or root.val == q:
self.res = branch
return branch + 1