Amount of Time for Binary Tree to Be Infected

Do exactly what the statement describes. First walk the tree once to index every node and record its parent, since the parent link is not given. Then let time advance one minute at a time: for every infected node, look at its two children and its parent, and collect the ones that are still healthy. Mark them infected at the end of the minute, and repeat until a minute infects nobody new. The number of minutes that passed is the answer.

from collections import deque


class Solution:
    def amountOfTime(self, root, start):
        # One pass to index every node and record its parent edge.
        nodes = {root.val: root}
        parent = {root.val: None}
        stack = [root]
        while stack:
            node = stack.pop()
            for child in (node.left, node.right):
                if child:
                    nodes[child.val] = child
                    parent[child.val] = node
                    stack.append(child)

        # Let time advance until a full minute infects nobody new.
        infected = {start}
        minutes = 0
        while len(infected) < len(nodes):
            newly = set()
            for val in infected:
                node = nodes[val]
                for other in (node.left, node.right, parent[val]):
                    if other and other.val not in infected:
                        newly.add(other.val)
            if not newly:
                break
            infected |= newly
            minutes += 1
        return minutes

Why it is quadratic: a correct simulation can need up to O(n) minutes — a path-shaped tree where start is one end infects a single new node per minute. Each minute rescans every currently infected node, so the total work is O(n) + O(n) + … over O(n) minutes, which is O(n²). The node and parent dictionaries cost O(n) space.

Stop advancing time explicitly and note what the simulation is really computing: the number of minutes a node waits is exactly its shortest-path distance from start in the tree treated as an undirected graph. So build that graph once. A single traversal adds, for every parent-child pair, both directions of the edge — the upward child -> parent edge is the one a binary tree is missing, and adding it is the entire trick. Then one BFS from start assigns every node its distance, and the largest distance assigned is the answer.

from collections import deque


class Solution:
    def amountOfTime(self, root, start):
        # Undirected adjacency: every node also lists its parent.
        adj = {}
        stack = [root]
        while stack:
            node = stack.pop()
            adj.setdefault(node.val, [])
            for child in (node.left, node.right):
                if child:
                    adj.setdefault(child.val, [])
                    adj[node.val].append(child.val)
                    adj[child.val].append(node.val)
                    stack.append(child)

        # BFS from start: each level is one minute, so the last one is the answer.
        dist = {start: 0}
        queue = deque([start])
        while queue:
            node = queue.popleft()
            for nxt in adj[node]:
                if nxt not in dist:
                    dist[nxt] = dist[node] + 1
                    queue.append(nxt)
        return max(dist.values())

Trace the first example, root = [1,5,3,null,4,10,6,9,2], start = 3. Watch the wave spread: each BFS level is one minute, and a node is reached just as easily through a child as through its parent.

Why it is linear: the tree has exactly n - 1 edges, so the adjacency list stores 2(n - 1) entries and the build is O(n). Every node enters the BFS queue once and is dequeued once, and each edge is examined a constant number of times, so the BFS is O(n) as well. Both structures hold O(n) values, and the BFS depth is O(n) in the worst case (a path), which is what the dist map already stores. The answer is the maximum recorded distance, matching Example 2 where the single node is at distance 0.