Problem 2385

(Last updated: )

Question statement:

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. Return the number of minutes needed for the entire tree to be infected, given that - each minute, a node becomes infected if:

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

Solution

Main idea:

The first thing to note is that the infection can spread from a child to a parent node. Thus, the tree structure is inherently a hinderence to solving the problem. Thus, our first step is to reinterpret the tree as an undirected graph. We accomplish this by going through the tree (dfs style) and creating an adjacency dictionary. Once we have the graph in place, we can proceed similar to a graph search approach.

We will have three sets:

  • cur_gen that tracks the the nodes that are infected when time t=k (initialized with the start node).
  • next_gen that tracks the nodes adjacent to those that belong to cur_gen (initialized to the empty array)
  • visited (or affected) that tracks the nodes that have already been visited.

The algorithm proceeds as follows:

  • Initialize a time counter to 0.
  • For each node in the current generation (cur_gen):
    • Scan its neighbors.
      • If a neighbor has been previously visited, continue
      • If a neighbor has not been previously visited, add the neighbor to next_gen and visited.
    • After scanning all neighbors, if we have a non-empty next_gen array, then add 1 to the time counter.
    • Delete this node from our graph (since we have dealt with it).
  • Update generation by setting cur_gen = next_gen and next_gen = []
  • Continue this way until the graph is empty.
  • Return time.

Code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
    def conv_to_graph(graph, root, parent_val):
        if root == None:
            return
        if(root.left != None):
                graph[root.val].append(root.left.val)
        if(root.right != None):
                graph[root.val].append(root.right.val)
        if(parent_val != 0):
            graph[root.val].append(parent_val)
        
        conv_to_graph(graph, root.left, root.val)
        conv_to_graph(graph, root.right, root.val)

        return graph
    
    gr = conv_to_graph(defaultdict(list), root, 0)
    cur_gen = set([start])
    visited = set([start])
    time = 0
    while gr:
        next_gen = set()
        for j in cur_gen:
            for k in gr[j]:
                if(k in visited):
                    continue
                else:
                    next_gen.add(k)
                    visited.add(k)
            del(gr[j])
        if(next_gen != set()):
                time += 1
        cur_gen = next_gen
    
    return time

Comments:

Time complexity: $\texttt{O(n$^2$)}$ where n is the number of nodes in the graph.

The above code runs in 427ms beats 82.38% of users. It uses 68.75MB of memory, beating 58.26% of users.

Also, I was travelling all day, and I submitted my solution (first try) with 11 minutes to spare. The adrenaline was real!