1740. 找到⼆叉树中的距离

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


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52