1740. 找到⼆叉树中的距离

这道题是最近公共祖先问题(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


10-14 03:46