leetcode - 2385. Amount of Time for Binary Tree to Be Infected

Description

You are given the root of a binary tree with unique values, and an integer start. At minute 0, an infection starts from the node with value start.

Each minute, a node becomes infected if:

The node is currently uninfected.
The node is adjacent to an infected node.

Return the number of minutes needed for the entire tree to be infected.

Example 1:

Input: root = [1,5,3,null,4,10,6,9,2], start = 3
Output: 4
Explanation: The following nodes are infected during:

• Minute 0: Node 3

• Minute 1: Nodes 1, 10 and 6

• Minute 2: Node 5

• Minute 3: Node 4

• Minute 4: Nodes 9 and 2
  It takes 4 minutes for the whole tree to be infected so we return 4.
  Example 2:
  

  Input: root = [1], start = 1
  Output: 0
  Explanation: At minute 0, the only node in the tree is infected so we return 0.

Constraints:

The number of nodes in the tree is in the range [1, 10^5].
1 <= Node.val <= 10^5
Each node has a unique value.
A node with a value of start exists in the tree.

Solution

Solved after hints, at first I thought I could use math and level traversal to solve this, but it turned out to be easier if transform the tree into a graph.

Hint: Transform the tree into a non-directional graph, then the answer is the longest road in the graph.

Time complexity: $o(n)$
Space complexity: $o(n)$

Code

# 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:
        graph = {}
        # dfs to build graph
        stack = [root]
        while stack:
            node = stack.pop()
            if node.val not in graph:
                graph[node.val] = []
            if node.right:
                graph[node.val].append(node.right.val)
                if node.right.val not in graph:
                    graph[node.right.val] = []
                graph[node.right.val].append(node.val)
                stack.append(node.right)
            if node.left:
                graph[node.val].append(node.left.val)
                if node.left.val not in graph:
                    graph[node.left.val] = []
                graph[node.left.val].append(node.val)
                stack.append(node.left)
        # bfs to get the farest node
        queue = collections.deque([(start, 0)])
        res = 0
        visited = set()
        while queue:
            node, step = queue.popleft()
            if node in visited:
                continue
            visited.add(node)
            res = max(res, step)
            for neighbor in graph[node]:
                queue.append((neighbor, step + 1))
        return res