Amount of Time for Binary Tree to Be Infected

• Step 1: Build Graph Representation.

  1. Create an adjacency list, for example, a Map<Integer, List<Integer>>, to store the graph structure.
  2. Traverse the binary tree using a recursive helper function (e.g., DFS).
  3. For each node, if it has a parent, add an undirected edge between the node and its parent in the adjacency list.
  4. Similarly, for each node, add undirected edges to its left and right children if they exist. A simpler way is to just traverse downwards and for each node with a child, add (node.val, child.val) and (child.val, node.val) to the map.

• Step 2: Perform Breadth-First Search (BFS).

  1. Initialize a queue and add the start node's value.
  2. Create a visited set to keep track of infected nodes, and add the start value to it.
  3. Initialize a variable minutes to -1.
  4. Begin a loop that continues as long as the queue is not empty. In each iteration, this loop processes one level of the graph (one minute of infection spread).
  5. Increment minutes.
  6. Get the number of nodes at the current level (levelSize).
  7. Loop levelSize times to process each node at the current level.
  8. Dequeue a node. For this node, iterate through its neighbors in the adjacency list.
  9. If a neighbor has not been visited, add it to the visited set and enqueue it.
  10. Once the main loop finishes, minutes will hold the maximum time required, which is the answer.

The core idea is to abstract away the tree structure into a graph, which simplifies reasoning about traversal in all directions (up, down, left, right).

First, we build the graph. We can traverse the tree with any method (like DFS). For each node, we establish connections to its neighbors. For any node u and its child v, we add an edge from u to v and another from v to u in our adjacency list. This ensures we can traverse both downwards and upwards.

1private void buildGraph(TreeNode node, Map<Integer, List<Integer>> adj) {2    if (node == null) return;3 4    if (node.left != null) {5        adj.computeIfAbsent(node.val, k -> new ArrayList<>()).add(node.left.val);6        adj.computeIfAbsent(node.left.val, k -> new ArrayList<>()).add(node.val);7    }8    if (node.right != null) {9        adj.computeIfAbsent(node.val, k -> new ArrayList<>()).add(node.right.val);10        adj.computeIfAbsent(node.right.val, k -> new ArrayList<>()).add(node.val);11    }12    buildGraph(node.left, adj);13    buildGraph(node.right, adj);14}

Once the graph is constructed, we perform a standard level-order BFS starting from the start node. BFS is ideal here because it explores the graph layer by layer, which directly corresponds to the spread of infection minute by minute. We keep track of the number of levels traversed; this count will be our final answer.

1public int amountOfTime(TreeNode root, int start) {2    Map<Integer, List<Integer>> adj = new HashMap<>();3    buildGraph(root, adj);4 5    Queue<Integer> queue = new LinkedList<>();6    queue.offer(start);7    Set<Integer> visited = new HashSet<>();8    visited.add(start);9    int minutes = -1;10 11    while (!queue.isEmpty()) {12        minutes++;13        int levelSize = queue.size();14        for (int i = 0; i < levelSize; i++) {15            int currentNode = queue.poll();16            if (adj.containsKey(currentNode)) {17                for (int neighbor : adj.get(currentNode)) {18                    if (!visited.contains(neighbor)) {19                        visited.add(neighbor);20                        queue.offer(neighbor);21                    }22                }23            }24        }25    }26    return minutes;27}

• Define a global or class-level variable maxTime and initialize it to 0.
• Create a recursive DFS function, say traverse(node, start), that returns an integer with special meaning:
  • A positive return value d signifies that the start node was found in the current node's subtree, and its distance from node is d.
  • A negative return value -h signifies that the start node was not found, and the height of the subtree rooted at node is h.
  • A return value of 0 corresponds to a null node.
• Base Case: If traverse is called on a null node, return 0.
• Recursive Step (Post-order Traversal):
  1. Recursively call traverse on the left and right children to get their results: leftResult and rightResult.
  2. Check if the current node is the start node. If node.val == start:
     • The infection starts here. The time to infect its own subtrees is the maximum of their heights. Update maxTime = max(maxTime, abs(leftResult), abs(rightResult)).
     • Return 1 to its parent, indicating the start node is at distance 1.
  3. If the start node was found in a child's subtree (e.g., leftResult > 0):
     • The infection path to the other subtree (right) goes up to the current node and then down. The time for this path is leftResult (up) + abs(rightResult) (down).
     • Update maxTime = max(maxTime, leftResult + abs(rightResult)).
     • Return leftResult + 1 to propagate the distance to start upwards.
     • Handle the symmetric case for the right subtree.
  4. If the start node is not in the current node's subtree (leftResult <= 0 and rightResult <= 0):
     • The function's job is to report the height of this subtree. Return -(1 + max(abs(leftResult), abs(rightResult))).

We perform a post-order traversal. For any given node, after visiting its children, we have two key pieces of information: the result from the left subtree and the result from the right subtree. Based on these results and the current node's value, we can determine the maximum infection time involving this node and its subtrees.

The recursive function traverse(node, start) cleverly encodes its findings in its return value:

• Positive value d: The start node is in this subtree, d edges away from the current node.
• Negative value -h: The start node is not in this subtree, and its height is h.

Let's analyze the logic at a node after recursive calls to its children:

1. If node is the start node: The infection spreads downwards. The time is the height of the taller child subtree. We update our global maxTime with this value. We then return 1 to the parent, signaling that the start node is its immediate child.
2. If start is in a subtree (e.g., left): We know the distance from start to the current node (which is leftResult). The longest path from start could be one that goes up to node and then all the way down the right subtree. The time for this path is leftResult + height_of_right_subtree. We update maxTime with this potential maximum. We then return leftResult + 1 to the parent, continuing to propagate the distance to start.
3. If start is not in the subtree: The node is just part of a regular branch. We calculate its height based on its children's heights and return it as a negative number to the parent.

This allows us to calculate both downward infection times (from start into its own subtree) and upward-then-downward infection times in a single traversal.

1class Solution {2    private int maxTime = 0;3 4    public int amountOfTime(TreeNode root, int start) {5        traverse(root, start);6        return maxTime;7    }8 9    /**10     * Returns the distance from node to start if start is in the subtree.11     * Otherwise, returns the negative height of the subtree.12     */13    private int traverse(TreeNode node, int start) {14        if (node == null) {15            return 0;16        }17 18        int leftResult = traverse(node.left, start);19        int rightResult = traverse(node.right, start);20 21        if (node.val == start) {22            // Infection starts here. Max time to infect its own subtrees.23            // The heights are the absolute values of the results from children.24            maxTime = Math.max(Math.abs(leftResult), Math.abs(rightResult));25            // Return 1 to parent, indicating distance to start.26            return 1;27        }28 29        if (leftResult > 0 || rightResult > 0) {30            // Start node is in one of the subtrees.31            // The infection has to travel up to the current node and then down the other subtree.32            maxTime = Math.max(maxTime, Math.abs(leftResult) + Math.abs(rightResult));33            // Propagate the distance to the start node upwards.34            return (leftResult > 0 ? leftResult : rightResult) + 1;35        } else {36            // Start node is not in this subtree. Return the height of this subtree as a negative value.37            return Math.min(leftResult, rightResult) - 1;38        }39    }40}