Amount of Time for Binary Tree to Be Infected
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