力扣 2385. 感染二叉树需要的总时间

题目来源：https://leetcode.cn/problems/amount-of-time-for-binary-tree-to-be-infected/

大致题意：
给一个二叉树和一个节点，假设该节点每单位时间可以感染与其相邻的节点，被感染的节点之后也会感染相邻节点，求整颗二叉树被感染需要的时间。
相邻节点是指某个节点的子节点和父节点

思路

比较简单的方法是：

• 遍历二叉树，存每个节点的父节点，构建一张图，然后使用给定节点作为源点，遍历求整张图其他节点距离源点的最短距离即可

但是该方法需要额外的空间来存储父节点，可以分析该问题得到该二叉树被感染的时间的计算方式：

1. 若当前节点是目标节点，那么以该节点为根节点的子树被感染需要的时间为 max(左子树被感染的时间, 右子树被感染的时间) + 1。这里子树被感染的时间即为子树的深度
2. 若目标节点在当前节点左子树，那么以该节点为根节点的右子树被感染需要的时间为 目标节点到当前节点的时间 + 右子树被感染的时间 + 1
3. 同理，若目标节点在当前节点右子树，那么以该节点为根节点的左子树被感染需要的时间为 目标节点到当前节点的时间 + 左子树被感染的时间 + 1
4. 还有一种情况，就是目标节点不在当前子树中，只是和当前子树有公共祖先，那么如何求当前子树被感染时间？其实对于这种情况，其子树被感染时间为 目标节点和当前节点之间的距离 + 当前子树的深度，而 目标节点和当前节点之间的距离 即为 当前节点和目标节点最近公共祖先与当前节点的距离 + 最近公共祖先与目标节点的距离，于是子树被感染的时间即为 最近公共祖先与当前节点的距离 + 最近公共祖先与目标节点的距离 + 当前子树的深度，同时 最近公共祖先与当前节点的距离 + 当前子树的深度 即为 公共祖先子树的深度，而这个值在遍历到公共祖先节点的时候已经算过了

在未找到目标节点时，可以认为某个节点感染其子节点的时间即为距离

于是可以通过 bfs 遍历求得上述所有感染时间的最大值即为答案，具体搜索时可以使用一个标志位存下目标节点所在层，同时每次遍历是标记当前节点所在层，这样当找到目标节点时，即可通过层数之间的差求得目标节点和祖先节点之间的感染时间。对于目标节点不在当前节点子树的情况，因为有标志位判断的情况存在，若此时目标节点已找到，相等于目标节点在左子树的情况（但是当前结果不会对答案有贡献，只是免去了特殊处理），否则不做特殊处理。

具体看代码：

class Solution {
    int ans;	// 二叉树被感染的时间
    int startDepth;	// 目标节点所在层
    int start;	// 目标节点值
    public int amountOfTime(TreeNode root, int start) {
        ans = 0;
        startDepth = -1;	// 初始化为 -1，表示未找到目标节点
        this.start = start;
        dfs(root, 0);
        return ans;
    }

    public int dfs(TreeNode node, int depth) {
        if (node == null) {
            return 0;
        }
        // 当前节点左子树的最大深度
        int l = dfs(node.left, depth + 1);
        // 如果此时目标节点已找到，说明目标节点在左子树中
        boolean inLeft = startDepth != -1;
        // 当前节点右子树的最大深度
        int r = dfs(node.right, depth + 1);
        // 如果此时目标节点找到，说明目标节点在右子树中
        boolean inRight = !inLeft && startDepth != -1;
        // 若当前节点为目标节点，更新目标节点深度和答案
        if (node.val == start) {
            startDepth = depth;
            ans = Math.max(Math.max(l, r), ans);
        } else if (inLeft) {	// 在左子树中，更新答案
            ans = Math.max(startDepth - depth + r, ans);
        } else if (inRight) {	// 在右子树中，更新答案
            ans = Math.max(startDepth - depth + l, ans);
        }
        return Math.max(l, r) + 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