题目 #
给你一棵二叉树的根节点 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