感染二叉树需要的总时间

难度: Medium

给你一棵二叉树的根节点 root ，二叉树中节点的值 互不相同 。另给你一个整数 start 。在第 0 分钟，感染将会从值为 start 的节点开始爆发。

每分钟，如果节点满足以下全部条件，就会被感染：

节点此前还没有感染。
节点与一个已感染节点相邻。

返回感染整棵树需要的分钟数。

示例 1：

输入：root = [1,5,3,null,4,10,6,9,2], start = 3
输出：4
解释：节点按以下过程被感染：
- 第 0 分钟：节点 3
- 第 1 分钟：节点 1、10、6
- 第 2 分钟：节点5
- 第 3 分钟：节点 4
- 第 4 分钟：节点 9 和 2
感染整棵树需要 4 分钟，所以返回 4 。

示例 2：

输入：root = [1], start = 1
输出：0
解释：第 0 分钟，树中唯一一个节点处于感染状态，返回 0 。

提示：

树中节点的数目在范围 [1, 10⁵] 内
1 <= Node.val <= 10⁵
每个节点的值 互不相同
树中必定存在值为 start 的节点

Submission

运行时间: 244 ms

内存: 0.0 MB

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
        self.traverse(root, start)
        return self.maxDistance
    
    def traverse(self, root, start):
        depth = 0
        if root is None:
            return depth
        
        leftDepth = self.traverse(root.left, start)
        rightDepth = self.traverse(root.right, start)

        if root.val == start:
            self.maxDistance = max(leftDepth, rightDepth)
            depth = -1
        elif leftDepth >= 0 and rightDepth >= 0:
            depth = max(leftDepth, rightDepth) + 1
        else:
            distance = abs(leftDepth) + abs(rightDepth)
            self.maxDistance = max(self.maxDistance, distance)
            depth = min(leftDepth, rightDepth) - 1
        
        return depth

Explain

这道题目是通过递归遍历二叉树来计算从特定节点开始感染整棵树所需的时间。首先，需要构建一个辅助函数 traverse 来深度优先搜索（DFS）整棵树，并记录从起始感染节点传播到树的最远点所需的时间。在遍历过程中，我们需要确定每个节点作为感染源时感染其子树的最大深度。如果当前节点是感染起始节点，我们就记录从这个节点到其子树的最大深度；如果不是，我们则计算从当前节点到已感染节点的距离，并更新最大感染时间。这种方法确保了我们可以在单次遍历中计算出感染整棵树所需的时间。

时间复杂度: O(n)

空间复杂度: O(n)

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
        # 初始化最大距离
        self.maxDistance = 0
        # 开始遍历树
        self.traverse(root, start)
        # 返回最大距离，即感染所有节点所需的时间
        return self.maxDistance
    
    def traverse(self, root, start):
        # 如果当前节点为空，返回0
        if root is None:
            return 0
        
        # 递归计算左子树深度
        leftDepth = self.traverse(root.left, start)
        # 递归计算右子树深度
        rightDepth = self.traverse(root.right, start)

        # 如果当前节点是感染起始点
        if root.val == start:
            # 计算从起始点到最远子节点的最大距离
            self.maxDistance = max(leftDepth, rightDepth)
            # 返回一个特殊值表示这是起始点
            return -1
        elif leftDepth >= 0 and rightDepth >= 0:
            # 如果当前节点不是起始点，更新深度
            return max(leftDepth, rightDepth) + 1
        else:
            # 如果已到达感染节点，计算和更新可能的最大距离
            distance = abs(leftDepth) + abs(rightDepth)
            self.maxDistance = max(self.maxDistance, distance)
            # 继续向上传递感染点信息
            return min(leftDepth, rightDepth) - 1

Explore

在`traverse`函数中，返回-1用于特殊标记感染起始点的节点。这是因为函数需要区分哪一个节点是感染的源头，从而能正确处理从源头向其他部分的感染传播。返回-1不仅帮助标记起始点，同时在递归过程中，使得起始点的父节点能够识别到从起始点开始的感染，并正确计算其他节点到起始点的距离。这种标记方式简化了从起始点到其他节点传播时间的计算过程。

当`leftDepth`和`rightDepth`均大于等于0时，这意味着当前节点不是感染起始点，并且它的子节点都已经被递归地探索过。这里增加并返回1的操作表示当前节点到其子节点的一步距离。因此，当我们从子节点返回深度时，需要在此基础上加1，以包括当前节点到子节点的这一层距离。这样，每返回一次，都会累加树中的一层，确保能够正确计算出从感染起始点到树中任意点的最大距离。

在该解法中，使用绝对值确实意味着`leftDepth`或`rightDepth`可能是负值。这里负值的使用是为了特别标记从感染起始点向上回溯的路径。对于从感染起始点开始的点，根据代码逻辑，会返回-1。因此，当我们在计算两个子树之间的距离时，可能会遇到包含负值的情况，这表示其中一个子树包含了感染起始点。使用绝对值的目的是确保无论感染起始点在左子树还是右子树，计算的距离总是正的，从而准确地反映从一侧的最远点到另一侧最远点的实际距离。