medium Amount of Time for Binary Tree to Be Infected

Problem Statement

You are given the root of a binary tree with unique values and an integer start.

At the start time (minute 0), an infection starts from the node with the value start.

Each minute, a node becomes infected if:

• It is currently uninfected.
• It is adjacent to an infected node.

Return the number of minutes required for the entire tree to become infected.

Examples

Example 1:

• Input: root = [1, 2, 3, null, 4, 5, 6], start = 3
• Expected Output: 3
• Explanation: The tree structure is:

      1
     / \
    2   3
    \  / \
     4 5  6
  

  • Minute 0: Node 3 is infected.
  • Minute 1: Nodes 1, 5, and 6 become infected.
  • Minute 2: Node 2 becomes infected.
  • Minute 3: Node 4 becomes infected.
  • Total time: 3 minutes.

Example 2:

• Input: root = [1, 2, 3, 4, 5, null, null, 6, 7, null, null, null, null, 8, 9], start = 5
• Expected Output: 4
• Explanation: The tree structure is:

           1
         /   \
        2     3
       / \   
      4   5
     / \  
    6   7
   / \
  8   9
  

  • Minute 0: Node 5 is infected.
  • Minute 1: Nodes 2 become infected.
  • Minute 2: Nodes 1, and 4 become infected.
  • Minute 3: Nodes 3, 6 and 7 become infected.
  • Minute 4: Nodes 8 and 9 become infected.
  • Total time: 4 minutes.

Example 3:

• Input: root = [10, 1, 2, null, 3, 4, 5], start = 1
• Expected Output: 3
• Explanation: The tree structure is:

      10
     /  \
    1    2
     \  / \
      3 4  5
  

  • Minute 0: Node 1 is infected.
  • Minute 1: Nodes 10 and 3 become infected.
  • Minute 2: Nodes 2 become infected.
  • Minute 3: Nodes 4 and 5 become infected.
  • Total time: 3 minutes.

Constraints:

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

Solution

To solve this problem, we can use a Breadth-First Search (BFS) approach starting from the node with the value equal to start. The BFS approach is suitable here because it allows us to explore all nodes at the current depth level before moving on to nodes at the next depth level, which mimics the infection spreading process minute by minute. This approach ensures that we efficiently track the infection spread through the tree by processing nodes level-by-level and marking them as infected.

By using BFS, we can keep a queue to manage the nodes to be processed and a set to track infected nodes. This method ensures that each node is processed exactly once, maintaining an O(n) time complexity, where n is the number of nodes in the tree.

Step-by-step Algorithm

1. Initialize Graph Representation:

   • Create an empty map or dictionary graph to represent the tree as an adjacency list.

2. Build Graph from Tree:

   • Define a helper function buildGraph(node, parent, graph) that takes the current node, its parent, and the graph as arguments.
   • In buildGraph:
     • If node is null, return.
     • If parent is not null, add an edge between node.val and parent.val in both directions in the graph.
     • Recursively call buildGraph for node.left and node.right.

3. Initialize BFS Structures:

   • Initialize a queue queue and add the start node to it.
   • Initialize a set infected to keep track of infected nodes and add the start node to it.
   • Initialize a variable minutes to -1 to count the number of minutes.

4. Breadth-First Search (BFS) for Infection Spread:

   • While queue is not empty:
     • Increment minutes by 1.
     • For each node in the current level (use a loop with the size of the queue):
       • Dequeue a node from queue.
       • For each neighbor of the dequeued node in graph:
         • If the neighbor is not in infected:
           • Add the neighbor to infected.
           • Enqueue the neighbor to queue.

5. Return the Result:

   • After the BFS loop ends, return minutes.

Algorithm Walkthrough

Given the tree:

    1
   / \
  2   3
   \  / \
    4 5  6

And the start node: 3

1. Initialize Graph Representation:

   • graph = {}

2. Build Graph from Tree:

   • Call buildGraph with root node (1):
     • graph = {1: []}
   • Call buildGraph with node 2 (left child of 1):
     • graph = {1: [2], 2: [1]}
   • Call buildGraph with node 4 (right child of 2):
     • graph = {1: [2], 2: [1, 4], 4: [2]}
   • Call buildGraph with node 3 (right child of 1):
     • graph = {1: [2, 3], 2: [1, 4], 3: [1], 4: [2]}
   • Call buildGraph with node 5 (left child of 3):
     • graph = {1: [2, 3], 2: [1, 4], 3: [1, 5], 4: [2], 5: [3]}
   • Call buildGraph with node 6 (right child of 3):
     • graph = {1: [2, 3], 2: [1, 4], 3: [1, 5, 6], 4: [2], 5: [3], 6: [3]}

3. Initialize BFS Structures:

   • queue = [3]
   • infected = {3}
   • minutes = -1

4. BFS for Infection Spread:

   • Minute 0:

     • minutes = 0
     • size = 1
     • Process node 3:
       • Dequeue 3
       • Process neighbor 1: Add 1 to infected and queue
       • Process neighbor 5: Add 5 to infected and queue
       • Process neighbor 6: Add 6 to infected and queue
     • queue = [1, 5, 6]
     • infected = {1, 3, 5, 6}

   • Minute 1:

     • minutes = 1
     • size = 3
     • Process node 1:
       • Dequeue 1
       • Process neighbor 2: Add 2 to infected and queue
       • Process neighbor 3: Already infected
     • Process node 5:
       • Dequeue 5
       • Process neighbor 3: Already infected
     • Process node 6:
       • Dequeue 6
       • Process neighbor 3: Already infected
     • queue = [2]
     • infected = {1, 2, 3, 5, 6}

   • Minute 2:

     • minutes = 2
     • size = 1
     • Process node 2:
       • Dequeue 2
       • Process neighbor 1: Already infected
       • Process neighbor 4: Add 4 to infected and queue
     • queue = [4]
     • infected = {1, 2, 3, 4, 5, 6}

   • Minute 3:

     • minutes = 3
     • size = 1
     • Process node 4:
       • Dequeue 4
       • Process neighbor 2: Already infected
     • queue = []
     • infected = {1, 2, 3, 4, 5, 6}

   • End of BFS

5. Return the Result:

   • return 3

Code

import java.util.*;

// class TreeNode {
//      int val;
//      TreeNode left;
//      TreeNode right;
//      TreeNode(int x) { val = x; }
//      TreeNode(int x, TreeNode left, TreeNode right) {
//          val = x;
//          this.left = left;
//         this.right = right;
//      }
// }

class Solution {

  public int amountOfTime(TreeNode root, int start) {
    // Build the graph representation of the tree
    Map<Integer, List<Integer>> graph = new HashMap<>();
    buildGraph(root, null, graph);

    // Queue for BFS
    Queue<Integer> queue = new LinkedList<>();
    // Set to keep track of infected nodes
    Set<Integer> infected = new HashSet<>();

    // Start the infection from the start node
    queue.add(start);
    infected.add(start);

    // Variable to count minutes
    int minutes = -1;

    // BFS to spread the infection
    while (!queue.isEmpty()) {
      int size = queue.size();
      minutes++;
      // Process all nodes at the current level
      for (int i = 0; i < size; i++) {
        int node = queue.poll();
        // Infect all neighbors
        for (int neighbor : graph.get(node)) {
          if (!infected.contains(neighbor)) {
            infected.add(neighbor);
            queue.add(neighbor);
          }
        }
      }
    }
    return minutes;
  }

  // Helper method to build the graph from the tree
  private void buildGraph(
    TreeNode node,
    TreeNode parent,
    Map<Integer, List<Integer>> graph
  ) {
    if (node == null) return;
    graph.putIfAbsent(node.val, new ArrayList<>());
    if (parent != null) {
      graph.get(node.val).add(parent.val);
      graph.get(parent.val).add(node.val);
    }
    buildGraph(node.left, node, graph);
    buildGraph(node.right, node, graph);
  }

  public static void main(String[] args) {
    Solution sol = new Solution();

    // Example 1
    TreeNode root1 = new TreeNode(
      1,
      new TreeNode(2, null, new TreeNode(4)),
      new TreeNode(3, new TreeNode(5), new TreeNode(6))
    );
    System.out.println(sol.amountOfTime(root1, 3)); // Output: 3

    // Example 2
    TreeNode root2 = new TreeNode(
      1,
      new TreeNode(
        2,
        new TreeNode(
          4,
          new TreeNode(6, new TreeNode(8), new TreeNode(9)),
          new TreeNode(7)
        ),
        null
      ),
      new TreeNode(3)
    );
    root2.left.right = new TreeNode(5);
    System.out.println(sol.amountOfTime(root2, 5)); // Output: 4

    // Example 3
    TreeNode root3 = new TreeNode(
      10,
      new TreeNode(1, null, new TreeNode(3)),
      new TreeNode(2, new TreeNode(4), new TreeNode(5))
    );
    System.out.println(sol.amountOfTime(root3, 1)); // Output: 3
  }
}

Complexity Analysis

Time Complexity

The time complexity for the solution is , where n is the number of nodes in the tree. This is because each node is processed exactly once during the BFS traversal.

Space Complexity

The space complexity is also . This is due to the space needed to store the graph representation of the tree and the additional space used by the BFS queue and the set of infected nodes.