Amount of Time for Binary Tree to Be Infected

Sanjeev SharmaSanjeev Sharma
2 min read

Advertisement

Problem

Starting from node start, infection spreads to adjacent nodes each minute. Return total minutes to infect the entire tree.

Key insight: Convert tree to undirected graph (add parent pointers), then BFS from start node. Answer = BFS depth.

Solutions

// C++
int amountOfTime(TreeNode* root, int start) {
    unordered_map<int, vector<int>> graph;
    function<void(TreeNode*, TreeNode*)> build = [&](TreeNode* node, TreeNode* par) {
        if (!node) return;
        if (par) { graph[node->val].push_back(par->val); graph[par->val].push_back(node->val); }
        build(node->left, node); build(node->right, node);
    };
    build(root, nullptr);
    unordered_set<int> visited;
    queue<int> q;
    q.push(start); visited.insert(start);
    int time = -1;
    while (!q.empty()) {
        time++;
        int sz = q.size();
        while (sz--) {
            int u = q.front(); q.pop();
            for (int v : graph[u]) if (!visited.count(v)) { visited.insert(v); q.push(v); }
        }
    }
    return time;
}
# Python
from collections import defaultdict, deque

def amountOfTime(root, start):
    graph = defaultdict(list)
    def build(node, parent):
        if not node:
            return
        if parent:
            graph[node.val].append(parent.val)
            graph[parent.val].append(node.val)
        build(node.left, node)
        build(node.right, node)
    build(root, None)
    visited = {start}
    q = deque([start])
    time = -1
    while q:
        time += 1
        for _ in range(len(q)):
            u = q.popleft()
            for v in graph[u]:
                if v not in visited:
                    visited.add(v)
                    q.append(v)
    return time
// Java
Map<Integer, List<Integer>> graph = new HashMap<>();
public int amountOfTime(TreeNode root, int start) {
    build(root, -1);
    Set<Integer> visited = new HashSet<>();
    Queue<Integer> q = new LinkedList<>();
    q.add(start); visited.add(start);
    int time = -1;
    while (!q.isEmpty()) {
        time++;
        int sz = q.size();
        while (sz-- > 0) {
            int u = q.poll();
            for (int v : graph.getOrDefault(u, new ArrayList<>()))
                if (!visited.contains(v)) { visited.add(v); q.add(v); }
        }
    }
    return time;
}
void build(TreeNode node, int parent) {
    if (node == null) return;
    graph.computeIfAbsent(node.val, k -> new ArrayList<>());
    if (parent != -1) { graph.get(node.val).add(parent); graph.get(parent).add(node.val); }
    build(node.left, node.val); build(node.right, node.val);
}
// JavaScript
function amountOfTime(root, start) {
    const graph = new Map();
    function build(node, parent) {
        if (!node) return;
        if (!graph.has(node.val)) graph.set(node.val, []);
        if (parent !== null) {
            graph.get(node.val).push(parent.val);
            if (!graph.has(parent.val)) graph.set(parent.val, []);
            graph.get(parent.val).push(node.val);
        }
        build(node.left, node); build(node.right, node);
    }
    build(root, null);
    const visited = new Set([start]);
    let queue = [start], time = -1;
    while (queue.length) {
        time++;
        const next = [];
        for (const u of queue)
            for (const v of (graph.get(u) || []))
                if (!visited.has(v)) { visited.add(v); next.push(v); }
        queue = next;
    }
    return time;
}

Complexity

  • Time: O(n)
  • Space: O(n)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro