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

题目

给你一棵二叉树的根节点 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, 105] 内

1 <= Node.val <= 105

每个节点的值 互不相同

树中必定存在值为 start 的节点

解法

我们使用后序遍历，对于每个节点维护以下信息：

该节点被感染的时间(尚未感染则为 -1)

以该节点为根的树的深度

对于当前遍历的节点，我们分以下情况考虑：

感染源就是该节点

感染以本节点为根的树所需时间 max(左子树深度, 右子树深度)

感染源在左子树中

感染以本节点为根的树所需时间

感染源在右子树中

感染源既不是该节点，也不在其左右子树中

代码


class Solution:
    def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
        ans = 0

        def dfs(node):
            nonlocal ans
            if not node:
                return (-1, 0)
            (leftInfectTime, leftDepth) = dfs(node.left)
            (rightInfectTime, rightDepth) = dfs(node.right)
            if node.val == start:
                ans = max(ans, max(leftDepth, rightDepth))
                return (0, 0)
            elif leftInfectTime >= 0:
                ans = max(ans, leftInfectTime + 1 + rightDepth)
                return (leftInfectTime + 1, 0)
            elif rightInfectTime >= 0:
                ans = max(ans, rightInfectTime + 1 + leftDepth)
                return (rightInfectTime + 1, 0)
            else:
                return (-1, max(leftDepth, rightDepth) + 1)
        
        dfs(root)
        return ans