发表于 2024.04.24

• 有点难度的二叉树递归，节点的感染路径有三个方向：往下感染左子树，右子树以及往上感染祖先。可以通过一次遍历得到结果：使用一个变量ans维护当前的感染最长路径，按后序遍历树节点，如果遇到了start节点，则先将ans设置为其子树的高度。否则，对于某个节点，如果其左子树包含start节点，则可能的最长感染路径为start节点->...[从左子树往上感染]->本节点->...[往下感染]->右子树的最深节点，右子树亦然，不断更新ans，直至处理到root节点。

  class Solution:
      def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
          ans = 0
  
          def dfs(node: Optional[TreeNode]) -> (int, int):
              # 返回值: 该子树的高, 该子树的根距离start节点的距离(如果start节点不在该子树内，则为-1)
              nonlocal ans
              if node is None:
                  return 0, -1
  
              left_depth, source_dis_l = dfs(node.left)
              right_depth, source_dis_r = dfs(node.right)
  
              if node.val == start:
                  # 碰到start节点
                  ans = max(left_depth, right_depth)
                  return max(left_depth, right_depth) + 1, 0
  
              if source_dis_l != -1:
                  # 左子树中包含start节点, 则可能的最长感染路径为 start节点 -> ... -> 本节点 -> 右子树最深节点
                  ans = max(ans, right_depth + 1 + source_dis_l)
              elif source_dis_r != -1:
                  # 右子树包含start节点, 则可能的最长感染路径为 start节点 -> ... -> 本节点 -> 左子树最深节点
                  ans = max(ans, left_depth + 1 + source_dis_r)
  
              # 计算本节点距离start节点的距离
              source_dis = max(source_dis_l, source_dis_r)
              if source_dis != -1:
                  source_dis += 1  # 加上本节点
              return 1 + max(left_depth, right_depth), source_dis
  
          dfs(root)
          return ans

• LC 题目链接