2385. Amount of Time for Binary Tree to Be Infected


Problem Description

In this problem, we are given the root of a binary tree where each node has a unique value, and an integer start which represents the value of the node from where an infection begins. At minute 0, the infection starts at the start node.

The infection spreads every minute to adjacent nodes that are not yet infected. A node is considered adjacent if it is directly connected to an infected node by an edge. The task is to calculate the number of minutes it will take for the infection to spread to the entire tree.

Intuition

The key to solving this problem is to consider the binary tree as an undirected graph, where an edge exists between parent and child nodes. We can then perform a breadth-first search (BFS) starting from the start node since BFS processes nodes level by level, which naturally aligns with the passage of minutes as the infection spreads.

The solution involves these steps:

  1. Convert the tree structure into a graph representation. This is necessary because a tree does not provide a straightforward way to move to siblings or the parent from a child node. In a graph representation, however, we can navigate to any adjacent node with ease.

  2. Once we have our graph, we begin a BFS from the start node. We keep a queue to keep track of nodes to process, and a set to remember which nodes we've already visited (to prevent reinfection).

  3. We process the nodes in the current queue, which represents all nodes that got infected in the current minute. For each of those nodes, we add their uninfected adjacent nodes to the queue to be processed in the next minute.

  4. The process continues until there are no more nodes to infect, which is when our queue is empty. We keep track of the time with a variable that increments at the end of processing all nodes at a current minute.

  5. The loop continues until the queue is empty, and the time variable at this point represents the total number of minutes taken to infect the entire tree.

The provided Python code follows this approach to arrive at the solution.

Learn more about Tree, Depth-First Search, Breadth-First Search and Binary Tree patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece๏ผš

Which of these pictures shows the visit order of a depth-first search?

Solution Approach

The given solution uses the Breadth-First Search (BFS) algorithm to simulate the spread of the infection across the tree. Here's a detailed breakdown of the implementation:

  1. Graph Construction:

    The solution begins by creating a graph representation of the tree using a default dictionary g of lists. This is done in the dfs (Depth-First Search) function.

    1def dfs(root):
    2    if root is None:
    3        return
    4    if root.left:
    5        g[root.val].append(root.left.val)
    6        g[root.left.val].append(root.val)
    7    if root.right:
    8        g[root.val].append(root.right.val)
    9        g[root.right.val].append(root.val)
    10    dfs(root.left)
    11    dfs(root.right)

    The dfs function is a recursive function that traverses the entire binary tree and adds each node's children to its list of adjacent nodes in the graph.

  2. Initialization:

    Initialize an empty set vis to keep track of visited nodes (infected nodes) and a queue q to maintain the BFS's order of node processing, with the start node as the initial node to be processed.

    1vis = set()
    2q = deque([start])
  3. BFS Algorithm:

    The solution sets up a while loop that continues until the queue q is empty, signifying that there are no more nodes to be infected.

    1ans = -1
    2while q:
    3    ans += 1
    4    for _ in range(len(q)):
    5        i = q.popleft()
    6        vis.add(i)
    7        for j in g[i]:
    8            if j not in vis:
    9                q.append(j)

    Inside the loop, ans is incremented to count the minutes. For each iteration of the while loop, it processes all nodes currently in the queue, which are the nodes that got infected in the previous minute. It pops each node from the queue, adds it to the vis set, and then iterates over its adjacent nodes. If any adjacent node has not been visited (infected), it is added to the queue to be processed in the next minute.

By the end of the BFS loop, ans holds the number of minutes it took to infect the entire tree since each loop iteration represents one minute of infection time passing.

The return value ans represents the required number of minutes for the infection to spread through the entire tree, as determined by our BFS algorithm.

1return ans

This solution efficiently converts the tree into a graph and then uses the BFS algorithm to simulate the spread of infection in discrete minute intervals.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Which algorithm is best for finding the shortest distance between two points in an unweighted graph?

Example Walkthrough

Let's illustrate the solution approach with a small example. Assume we have the following binary tree and the infection starts at the node with value 3.

1       1
2      / \
3     2   3
4        / \
5       4   5

Step 1: Convert the tree into a graph representation.

Following the DFS traversal, we get:

  • Node 1 is connected to nodes 2 and 3.
  • Node 2 is connected to node 1.
  • Node 3 is connected to nodes 1, 4, and 5.
  • Node 4 is connected to node 3.
  • Node 5 is connected to node 3.

Our graph g representation will be:

1{
2    1: [2, 3],
3    2: [1],
4    3: [1, 4, 5],
5    4: [3],
6    5: [3]
7}

Step 2: Initialize the visited set vis to {} and the queue q to [3] since infection starts at node 3.

Step 3: Begin BFS algorithm.

  • At minute 0, ans = -1. We increment ans to 0. The queue q has just one node, which is 3. Node 3 is popped from the queue and added to the visited set: vis = {3}. Nodes 1, 4, and 5 are adjacent to 3 and are added to the queue: q = [1, 4, 5].

  • At minute 1, ans is incremented to 1. Three nodes (1, 4, and 5) are in the queue. We visit each node, add it to the visited set, and add their unvisited (uninfected) adjacent nodes to the queue. However, node 1 has no unvisited adjacent nodes, and both nodes 4 and 5 have only node 3 as their adjacent node, which is already visited. So the queue remains empty after this: vis = {1, 3, 4, 5} and q = [].

There are no more nodes to infect and the queue is empty.

By the end of the process, ans = 1, which means that it took 1 minute for the infection to spread through the entire tree.

This small example illustrates how the BFS algorithm spreads the infection across the tree level by level, minutely fitting the problem's requirement to calculate the time it takes to infect the whole tree.

Solution Implementation

1# Importing the required modules.
2from collections import defaultdict, deque
3
4# Definition for a binary tree node.
5class TreeNode:
6    def __init__(self, val=0, left=None, right=None):
7        self.val = val
8        self.left = left
9        self.right = right
10
11class Solution:
12    def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
13        # Depth-First Search to build the graph from the binary tree.
14        def dfs(node):
15            if node is None:
16                return
17            # Connect current node with left child if it exists.
18            if node.left:
19                graph[node.val].append(node.left.val)
20                graph[node.left.val].append(node.val)
21            # Connect current node with right child if it exists.
22            if node.right:
23                graph[node.val].append(node.right.val)
24                graph[node.right.val].append(node.val)
25            # Recursive call for left and right children.
26            dfs(node.left)
27            dfs(node.right)
28
29        # Initialize a default dictionary for the graph representation.
30        graph = defaultdict(list)
31        # Start DFS traversal to build graph.
32        dfs(root)
33        # Initialize a set to keep track of visited nodes.
34        visited = set()
35        # Initialize a queue with the starting node.
36        queue = deque([start])
37        # Initialize time counter.
38        time = -1
39        # Execute until the queue is empty.
40        while queue:
41            # Increase time each level we go deeper in the graph.
42            time += 1
43            # Traverse nodes at the current level.
44            for _ in range(len(queue)):
45                current_node = queue.popleft()
46                visited.add(current_node)
47                # Explore all the neighbors of the current node.
48                for neighbor in graph[current_node]:
49                    if neighbor not in visited:
50                        queue.append(neighbor)
51        # Return the total amount of time.
52        return time
53
1class Solution {
2    // To store the adjacency list representation of the binary tree
3    private Map<Integer, List<Integer>> adjacencyList = new HashMap<>();
4
5    // Calculates the amount of time needed to infect all nodes starting from 'start' node
6    public int amountOfTime(TreeNode root, int start) {
7        // Convert tree to graph (adjacency list)
8        convertToGraph(root);
9        // Queue for the Breadth First Search (BFS)
10        Deque<Integer> queue = new ArrayDeque<>();
11        // Set to track visited nodes
12        Set<Integer> visited = new HashSet<>();
13      
14        // Offer the starting node into the queue
15        queue.offer(start);
16        int time = -1; // Initialize time to -1 because we increase time before processing the nodes
17      
18        while (!queue.isEmpty()) {
19            time++; // Increase time as we are moving to the next level of nodes
20          
21            // Loop through the nodes at the current level
22            for (int i = queue.size(); i > 0; i--) {
23                int currentNode = queue.pollFirst();
24                visited.add(currentNode);
25              
26                // Get the neighbors of the current node
27                if (adjacencyList.containsKey(currentNode)) {
28                    for (int neighbor : adjacencyList.get(currentNode)) {
29                        if (!visited.contains(neighbor)) {
30                            queue.offer(neighbor);
31                        }
32                    }
33                }
34            }
35        }
36        return time; // Return the total time taken to infect all nodes
37    }
38
39    // Helper method to convert the binary tree to a graph represented as an adjacency list
40    private void convertToGraph(TreeNode node) {
41        if (node == null) {
42            return;
43        }
44      
45        // Connect the current node with its left child
46        if (node.left != null) {
47            adjacencyList.computeIfAbsent(node.val, k -> new ArrayList<>()).add(node.left.val);
48            adjacencyList.computeIfAbsent(node.left.val, k -> new ArrayList<>()).add(node.val);
49        }
50      
51        // Connect the current node with its right child
52        if (node.right != null) {
53            adjacencyList.computeIfAbsent(node.val, k -> new ArrayList<>()).add(node.right.val);
54            adjacencyList.computeIfAbsent(node.right.val, k -> new ArrayList<>()).add(node.val);
55        }
56      
57        // Recursively convert the left and right subtrees to the graph
58        convertToGraph(node.left);
59        convertToGraph(node.right);
60    }
61}
62
1#include <unordered_map>
2#include <vector>
3#include <queue>
4#include <unordered_set>
5
6// Definition for a binary tree node
7struct TreeNode {
8    int val;
9    TreeNode *left;
10    TreeNode *right;
11    TreeNode() : val(0), left(nullptr), right(nullptr) {}
12    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
13    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
14};
15
16class Solution {
17public:
18    // Graph representation of the tree
19    unordered_map<int, vector<int>> graph;
20
21    // Find the amount of time needed to infect all nodes starting from the start node
22    int amountOfTime(TreeNode* root, int start) {
23        // Build graph representation of the tree using DFS
24        constructGraph(root);
25
26        // Queue for BFS traversal
27        queue<int> q;
28        q.push(start);
29
30        // Set to keep track of visited nodes
31        unordered_set<int> visited;
32
33        // Variable to track the number of minutes passed
34        int minutesPassed = -1;
35
36        // Perform BFS on the graph
37        while (!q.empty()) {
38            // Increment time for each level
39            ++minutesPassed;
40            // Iterate over all nodes in the current level
41            for (int levelSize = q.size(); levelSize > 0; --levelSize) {
42                int currentNode = q.front();
43                q.pop();
44                visited.insert(currentNode);
45                // Add all unvisited adjacent nodes to the queue
46                for (int adjacentNode : graph[currentNode]) {
47                    if (!visited.count(adjacentNode)) {
48                        q.push(adjacentNode);
49                    }
50                }
51            }
52        }
53
54        // Return the total time passed to infect all nodes
55        return minutesPassed;
56    }
57
58    // Helper function to build graph from the binary tree using DFS
59    void constructGraph(TreeNode* root) {
60        if (!root) return;
61
62        // Connect current node with its left child (if exists)
63        if (root->left) {
64            graph[root->val].push_back(root->left->val);
65            graph[root->left->val].push_back(root->val);
66        }
67
68        // Connect current node with its right child (if exists)
69        if (root->right) {
70            graph[root->val].push_back(root->right->val);
71            graph[root->right->val].push_back(root->val);
72        }
73
74        // Continue DFS with the children
75        constructGraph(root->left);
76        constructGraph(root->right);
77    }
78};
79
1/**
2 * Given the root of a binary tree and an integer `start` representing the value of the node where an
3 * infection starts, the function `amountOfTime` returns the minimum number of minutes needed to infect all nodes in the tree.
4 * Here, each minute, any infected node in the tree can spread the infection to its children and its parent.
5 *
6 * @param {TreeNode | null} root - The root node of the binary tree.
7 * @param {number} start - The value of the node where the infection starts.
8 * @returns The minimum number of minutes needed to infect all nodes.
9 */
10
11// Auxiliary function to perform Depth First Search (DFS) starting from a node while avoiding the parent node
12// to compute the time taken to spread the infection throughout the tree.
13function depthFirstSearch(start: number, parent: number): number {
14    let maxTime = 0;
15    // Retrieve adjacent nodes (children and parent) for the current node.
16    const adjacentNodes = adjacencyList.get(start) ?? [];
17    for (const nextNode of adjacentNodes) {
18        // Avoid returning to the parent node.
19        if (nextNode !== parent) {
20            // Recur on the adjacent node, incrementing the time by 1.
21            const time = depthFirstSearch(nextNode, start) + 1;
22            // Keep track of the maximum time found so far.
23            maxTime = Math.max(maxTime, time);
24        }
25    }
26    return maxTime;
27}
28
29// Function to convert the binary tree into an adjacency list, which will allow easy traversal of
30// the nodes in no particular order to help simulate the spread of the infection.
31function createAdjacencyList(node: TreeNode) {
32    if (node.left != null) {
33        // If there is a left child, connect parent to left and vice versa.
34        adjacencyList.set(node.val, [...(adjacencyList.get(node.val) ?? []), node.left.val]);
35        adjacencyList.set(node.left.val, [...(adjacencyList.get(node.left.val) ?? []), node.val]);
36        // Recur on the left child.
37        createAdjacencyList(node.left);
38    }
39    if (node.right != null) {
40        // If there is a right child, connect parent to right and vice versa.
41        adjacencyList.set(node.val, [...(adjacencyList.get(node.val) ?? []), node.right.val]);
42        adjacencyList.set(node.right.val, [...(adjacencyList.get(node.right.val) ?? []), node.val]);
43        // Recur on the right child.
44        createAdjacencyList(node.right);
45    }
46}
47
48// This map will store the adjacency list of the tree where the key is the node value
49// and the value is an array of adjacent node values.
50const adjacencyList = new Map<number, number[]>();
51
52// This is our main function which will first construct an adjacency list from the given tree
53// and then use DFS to find out the time taken to spread the infection from the 'start' node to all other nodes.
54function amountOfTime(root: TreeNode | null, start: number): number {
55    // Prepare the adjacency list from the binary tree.
56    if (root) createAdjacencyList(root);
57    // Perform DFS starting at the start node, using -1 to indicate that there's no parent for the start node.
58    return depthFirstSearch(start, -1);
59}
60
Not Sure What to Study? Take the 2-min Quiz๏ผš

Is the following code DFS or BFS?

1void search(Node root) {
2  if (!root) return;
3  visit(root);
4  root.visited = true;
5  for (Node node in root.adjacent) {
6    if (!node.visited) {
7      search(node);
8    }
9  }
10}

Time and Space Complexity

Time Complexity:

The time complexity of the provided code is O(N) where N is the number of nodes in the tree. The reason for this is as follows:

  • The dfs function is called once for each node in the tree during the construction of the graph g. dfs visits each node exactly once.
  • The BFS loop, which uses the queue q, runs while there are nodes to be visited. In the worst case, each node is inserted into and removed from the queue exactly once, which leads to O(N) operations.

Therefore, because both the DFS and the BFS visit each node exactly once, the overall time complexity is O(N).

Space Complexity:

The space complexity of the code is also O(N) for the following reasons:

  • The graph g, which is a defaultdict(list), stores an adjacency list representation of the tree. This requires space for each edge in the undirected graph representation of the tree. Since a tree with N nodes has N-1 edges and each edge is stored twice (once for each direction), the space required is 2*(N-1), which is O(N).
  • The set vis, which keeps track of visited nodes, stores up to N node values in the worst case.
  • The queue q may store up to N node values in the worst case when nodes are waiting to be visited.
  • The recursive stack space for dfs will have a maximum depth of the height of the tree, which is O(N) in the worst case for a skewed tree.

Summing these up, the space used is O(N) due to the storage of nodes in the graph representation, the visited set, the queue, and the call stack.

Learn more about how to find time and space complexity quickly using problem constraints.

Fast Track Your Learning with Our Quick Skills Quiz:

What's the output of running the following function using input [30, 20, 10, 100, 33, 12]?

1def fun(arr: List[int]) -> List[int]:
2    import heapq
3    heapq.heapify(arr)
4    res = []
5    for i in range(3):
6        res.append(heapq.heappop(arr))
7    return res
8
1public static int[] fun(int[] arr) {
2    int[] res = new int[3];
3    PriorityQueue<Integer> heap = new PriorityQueue<>();
4    for (int i = 0; i < arr.length; i++) {
5        heap.add(arr[i]);
6    }
7    for (int i = 0; i < 3; i++) {
8        res[i] = heap.poll();
9    }
10    return res;
11}
12
1class HeapItem {
2    constructor(item, priority = item) {
3        this.item = item;
4        this.priority = priority;
5    }
6}
7
8class MinHeap {
9    constructor() {
10        this.heap = [];
11    }
12
13    push(node) {
14        // insert the new node at the end of the heap array
15        this.heap.push(node);
16        // find the correct position for the new node
17        this.bubble_up();
18    }
19
20    bubble_up() {
21        let index = this.heap.length - 1;
22
23        while (index > 0) {
24            const element = this.heap[index];
25            const parentIndex = Math.floor((index - 1) / 2);
26            const parent = this.heap[parentIndex];
27
28            if (parent.priority <= element.priority) break;
29            // if the parent is bigger than the child then swap the parent and child
30            this.heap[index] = parent;
31            this.heap[parentIndex] = element;
32            index = parentIndex;
33        }
34    }
35
36    pop() {
37        const min = this.heap[0];
38        this.heap[0] = this.heap[this.size() - 1];
39        this.heap.pop();
40        this.bubble_down();
41        return min;
42    }
43
44    bubble_down() {
45        let index = 0;
46        let min = index;
47        const n = this.heap.length;
48
49        while (index < n) {
50            const left = 2 * index + 1;
51            const right = left + 1;
52
53            if (left < n && this.heap[left].priority < this.heap[min].priority) {
54                min = left;
55            }
56            if (right < n && this.heap[right].priority < this.heap[min].priority) {
57                min = right;
58            }
59            if (min === index) break;
60            [this.heap[min], this.heap[index]] = [this.heap[index], this.heap[min]];
61            index = min;
62        }
63    }
64
65    peek() {
66        return this.heap[0];
67    }
68
69    size() {
70        return this.heap.length;
71    }
72}
73
74function fun(arr) {
75    const heap = new MinHeap();
76    for (const x of arr) {
77        heap.push(new HeapItem(x));
78    }
79    const res = [];
80    for (let i = 0; i < 3; i++) {
81        res.push(heap.pop().item);
82    }
83    return res;
84}
85

Recommended Readings


Got a question?ย Ask the Teaching Assistantย anything you don't understand.

Still not clear? Ask in the Forum, ย Discordย orย Submitย the part you don't understand to our editors.

โ†
โ†‘TA ๐Ÿ‘จโ€๐Ÿซ