2385. 感染二叉樹需要的總時間

題目連結	2385. 感染二叉樹需要的總時間
思路	識別為“樹的直徑”
題解連結	從兩次遍歷到一次遍歷（Python/Java/C++/Go/JS/Rust）
關鍵點	1. 返回值為(感染時間, 是否包含感染源) 2. 分為三種情況：本節點為感染源、左兒子（或右兒子）為感染源、該子樹無感染源
時間複雜度	\(O(n)\)
空間複雜度	\(O(n)\)

程式碼實現：

class Solution:
    def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
        answer = 0
        def dfs(node):
            if node is None:
                return 0, False
            l_len, l_found = dfs(node.left)
            r_len, r_found = dfs(node.right)
            nonlocal answer
            if node.val == start:
                answer = max(l_len, r_len)
                return 1, True
            if l_found or r_found:
                answer = max(answer, l_len + r_len)
                return 1 + (l_len if l_found else r_len), True
            return max(l_len, r_len) + 1, False
        dfs(root)
        return answer