LeetCode 2385 - Amount of Time for Binary Tree to Be Infected

This problem asks us to model the spread of an infection across a binary tree. Each node in the tree has a unique integer value, and the infection starts from a specific node, identified by start.

leetcodemediumhash-tabletreedepth-first-searchbreadth-first-searchbinary-tree

LeetCode Problem 2385

Difficulty: 🟡 Medium
Topics: Hash Table, Tree, Depth-First Search, Breadth-First Search, Binary Tree

Solution

Problem Understanding

This problem asks us to model the spread of an infection across a binary tree. Each node in the tree has a unique integer value, and the infection starts from a specific node, identified by start. The infection spreads to adjacent nodes-this means both the parent and child nodes of an infected node-every minute. Our goal is to calculate the total number of minutes required until every node in the tree becomes infected.

The input is a binary tree represented by its root node, along with a starting node value. The output is a single integer representing the time taken to infect the entire tree. The constraints tell us the tree can be quite large, up to 100,000 nodes, and each node value is unique, so we do not need to worry about duplicate nodes. Every node is guaranteed to exist and the node with the starting value is present.

Important edge cases include a tree with a single node, where the infection is already complete at minute 0, and highly unbalanced trees, where the longest path from the start node might be deep on one side of the tree. A naive recursive DFS without considering the parent relationships could fail to propagate the infection upwards efficiently.

Approaches

A brute-force approach would involve repeatedly traversing the tree to find all uninfected neighbors at each minute. For each minute, we could do a DFS or BFS from the start node and infect neighbors incrementally. While this approach is correct, it is inefficient because each traversal could take O(n) time and we might repeat this for every minute until the tree is fully infected, resulting in O(n^2) time complexity for skewed trees.

The key observation that leads to an optimal solution is that the infection spreads as a wave, similar to BFS. However, standard tree traversal does not account for moving upwards from child to parent. To overcome this, we can preprocess the tree into a graph representation using a hash map that records adjacency for each node, including both children and parents. Once the adjacency map is built, a BFS starting from the start node efficiently models the spread of the infection minute by minute. This reduces the problem to a single BFS traversal over the entire tree.

Approach Time Complexity Space Complexity Notes
Brute Force O(n^2) O(n) Repeated DFS per minute, slow for deep or large trees
Optimal O(n) O(n) Convert tree to adjacency graph, BFS once

Algorithm Walkthrough

  1. Convert Tree to Graph: Traverse the tree once using DFS. For each node, record its children in an adjacency map, and also record the parent for upward traversal. This ensures we can traverse to any adjacent node efficiently.
  2. Initialize BFS: Use a queue initialized with the starting node. Also maintain a visited set to track infected nodes. Start a minutes counter at -1 since minute 0 will be counted during BFS.
  3. BFS Traversal: While the queue is not empty, increment the minutes counter. For each node in the current level of the BFS, mark its unvisited neighbors as visited and enqueue them. This step ensures infection spreads to all adjacent nodes in each minute.
  4. Termination: When the queue is empty, all nodes have been infected. Return the minutes counter, which now reflects the total number of minutes required to infect the entire tree.

Why it works: The BFS ensures that all nodes at distance d from the start node are infected at minute d. The adjacency map guarantees that both upward and downward propagation are considered. Therefore, the last BFS level corresponds to the maximum distance from the start node to any other node, giving the total infection time.

Python Solution

# 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
from collections import deque, defaultdict
from typing import Optional

class Solution:
    def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
        # Step 1: Build adjacency graph
        graph = defaultdict(list)
        
        def dfs(node, parent):
            if not node:
                return
            if parent:
                graph[node.val].append(parent.val)
                graph[parent.val].append(node.val)
            if node.left:
                dfs(node.left, node)
            if node.right:
                dfs(node.right, node)
        
        dfs(root, None)
        
        # Step 2: BFS to simulate infection
        queue = deque([start])
        visited = set([start])
        minutes = -1
        
        while queue:
            minutes += 1
            for _ in range(len(queue)):
                node = queue.popleft()
                for neighbor in graph[node]:
                    if neighbor not in visited:
                        visited.add(neighbor)
                        queue.append(neighbor)
        
        return minutes

Explanation: The DFS builds a bidirectional graph where each node points to its neighbors. BFS starts from the start node and spreads the infection one minute per level. The minutes counter increments at each level, and the traversal stops when all nodes have been visited.

Go Solution

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func amountOfTime(root *TreeNode, start int) int {
    // Build adjacency graph
    graph := make(map[int][]int)
    var dfs func(node, parent *TreeNode)
    dfs = func(node, parent *TreeNode) {
        if node == nil {
            return
        }
        if parent != nil {
            graph[node.Val] = append(graph[node.Val], parent.Val)
            graph[parent.Val] = append(graph[parent.Val], node.Val)
        }
        dfs(node.Left, node)
        dfs(node.Right, node)
    }
    dfs(root, nil)
    
    // BFS to simulate infection
    queue := []int{start}
    visited := make(map[int]bool)
    visited[start] = true
    minutes := -1
    
    for len(queue) > 0 {
        minutes++
        size := len(queue)
        for i := 0; i < size; i++ {
            node := queue[0]
            queue = queue[1:]
            for _, neighbor := range graph[node] {
                if !visited[neighbor] {
                    visited[neighbor] = true
                    queue = append(queue, neighbor)
                }
            }
        }
    }
    
    return minutes
}

Go-specific notes: We use slices to implement the queue and a map to track visited nodes. nil checks replace Python’s None, and slices dynamically grow with append. BFS logic mirrors Python but uses explicit indexing.

Worked Examples

Example 1:

Tree nodes: [1,5,3,null,4,10,6,9,2], start = 3

Adjacency graph after DFS:

1: [3]
3: [1, 10, 6]
5: [4]
4: [5, 9, 2]
10: [3]
6: [3]
9: [4]
2: [4]

BFS Levels:

Minute Queue Infected Nodes
0 [3] 3
1 [1, 10, 6] 1, 10, 6
2 [5] 5
3 [4] 4
4 [9, 2] 9, 2

Output: 4

Example 2:

Single node tree, start = 1. BFS queue starts with [1], no neighbors. Infection complete at minute 0.

Complexity Analysis

Measure Complexity Explanation
Time O(n) Each node is visited once in DFS to build the graph and once in BFS to simulate infection
Space O(n) Adjacency map stores each node and its neighbors, BFS queue and visited set may hold all nodes in the worst case

The linear time and space complexity is optimal given that every node must be processed at least once.

Test Cases

# Provided examples
assert Solution().amountOfTime(TreeNode(1), 1) == 0  # single node
root1 = TreeNode(1)
root1.left = TreeNode(5)
root1.right = TreeNode(3)
root1.left.right = TreeNode(4)
root1.left.right.left = TreeNode(9)
root1.left.right.right = TreeNode(2)
root1.right.left = TreeNode(10)
root1.right.right = TreeNode(6)
assert Solution().amountOfTime(root1, 3) == 4  # complex tree

# Edge cases
assert Solution().amountOfTime(TreeNode(1, TreeNode(2)), 2) == 1  # start at leaf
assert Solution().amountOfTime(TreeNode(1, None, TreeNode(2)), 1) == 1  # start at root, right child
assert Solution().amountOfTime(TreeNode(1), 1) == 0  # single node
Test Why
single node ensures function handles smallest tree
complex tree validates BFS spreads correctly and counts minutes accurately
start at leaf tests upward propagation
start at root with one child checks basic propagation to