Welcome to Subscribe On Youtube

2385. Amount of Time for Binary Tree to Be Infected

Description

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.

Each minute, a node becomes infected if:

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

Return the number of minutes needed for the entire tree to be infected.

 

Example 1:

Input: root = [1,5,3,null,4,10,6,9,2], start = 3
Output: 4
Explanation: The following nodes are infected during:
- Minute 0: Node 3
- Minute 1: Nodes 1, 10 and 6
- Minute 2: Node 5
- Minute 3: Node 4
- Minute 4: Nodes 9 and 2
It takes 4 minutes for the whole tree to be infected so we return 4.

Example 2:

Input: root = [1], start = 1
Output: 0
Explanation: At minute 0, the only node in the tree is infected so we return 0.

 

Constraints:

  • The number of nodes in the tree is in the range [1, 105].
  • 1 <= Node.val <= 105
  • Each node has a unique value.
  • A node with a value of start exists in the tree.

Solutions

Solution 2: Explicit Stack

Thinking

Infection spreads along parent and child edges, so the minutes needed are the farthest distance from $start$. Recording those edges and then returning a depth is correct on a short tree.

The tree can contain $10^5$ nodes. A left chain makes both walks recurse once per node, which exceeds the call stack.

Each walk only needs the current node and its parent, and a node’s depth is one plus the maximum depth of its other neighbors. Those values can sit beside an explicit stack.

The first stack records every parent edge in an undirected adjacency list, using the distinct node values as keys. The second stack enters a node, pushes its neighbors, and on exit writes one plus the maximum neighbor depth. The depth stored for $start$ is the answer.

We traverse the binary tree with an explicit stack and record each parent-child edge in an undirected adjacency list $g$. Node values are distinct, so they are the graph keys.

A second explicit stack then walks outward from $start$. On entry it pushes an exit marker and every neighbor except the parent; on exit the depth is one plus the maximum depth of those neighbors. The depth stored for $start$ is the answer.

The time complexity is $O(n)$, and the space complexity is $O(n)$, where $n$ is the number of nodes in the binary tree.

  • /**
     * Definition for a binary tree node.
     * public class TreeNode {
     *     int val;
     *     TreeNode left;
     *     TreeNode right;
     *     TreeNode() {}
     *     TreeNode(int val) { this.val = val; }
     *     TreeNode(int val, TreeNode left, TreeNode right) {
     *         this.val = val;
     *         this.left = left;
     *         this.right = right;
     *     }
     * }
     */
    class Solution {
        private Map<Integer, List<Integer>> g = new HashMap<>();
    
        public int amountOfTime(TreeNode root, int start) {
            dfs(root);
            Deque<Integer> q = new ArrayDeque<>();
            Set<Integer> vis = new HashSet<>();
            q.offer(start);
            int ans = -1;
            while (!q.isEmpty()) {
                ++ans;
                for (int n = q.size(); n > 0; --n) {
                    int i = q.pollFirst();
                    vis.add(i);
                    if (g.containsKey(i)) {
                        for (int j : g.get(i)) {
                            if (!vis.contains(j)) {
                                q.offer(j);
                            }
                        }
                    }
                }
            }
            return ans;
        }
    
        private void dfs(TreeNode root) {
            if (root == null) {
                return;
            }
            if (root.left != null) {
                g.computeIfAbsent(root.val, k -> new ArrayList<>()).add(root.left.val);
                g.computeIfAbsent(root.left.val, k -> new ArrayList<>()).add(root.val);
            }
            if (root.right != null) {
                g.computeIfAbsent(root.val, k -> new ArrayList<>()).add(root.right.val);
                g.computeIfAbsent(root.right.val, k -> new ArrayList<>()).add(root.val);
            }
            dfs(root.left);
            dfs(root.right);
        }
    }
    
    
    // Solution 2
    /**
     * Definition for a binary tree node.
     * public class TreeNode {
     *     int val;
     *     TreeNode left;
     *     TreeNode right;
     *     TreeNode() {}
     *     TreeNode(int val) { this.val = val; }
     *     TreeNode(int val, TreeNode left, TreeNode right) {
     *         this.val = val;
     *         this.left = left;
     *         this.right = right;
     *     }
     * }
     */
    class Solution {
        private static class Frame {
            TreeNode node;
            TreeNode fa;
    
            Frame(TreeNode node, TreeNode fa) {
                this.node = node;
                this.fa = fa;
            }
        }
    
        public int amountOfTime(TreeNode root, int start) {
            Map<Integer, List<Integer>> g = new HashMap<>();
            Deque<Frame> stk = new ArrayDeque<>();
            if (root != null) {
                stk.push(new Frame(root, null));
            }
            while (!stk.isEmpty()) {
                Frame cur = stk.pop();
                TreeNode node = cur.node;
                TreeNode fa = cur.fa;
                if (fa != null) {
                    g.computeIfAbsent(node.val, k -> new ArrayList<>()).add(fa.val);
                    g.computeIfAbsent(fa.val, k -> new ArrayList<>()).add(node.val);
                }
                if (node.right != null) {
                    stk.push(new Frame(node.right, node));
                }
                if (node.left != null) {
                    stk.push(new Frame(node.left, node));
                }
            }
            Map<Integer, Integer> dist = new HashMap<>();
            Deque<int[]> walk = new ArrayDeque<>();
            walk.push(new int[] {start, -1, 0});
            while (!walk.isEmpty()) {
                int[] cur = walk.pop();
                int node = cur[0], fa = cur[1], state = cur[2];
                List<Integer> nxts = g.getOrDefault(node, List.of());
                if (state == 0) {
                    walk.push(new int[] {node, fa, 1});
                    for (int i = nxts.size() - 1; i >= 0; --i) {
                        int nxt = nxts.get(i);
                        if (nxt != fa) {
                            walk.push(new int[] {nxt, node, 0});
                        }
                    }
                } else {
                    int best = 0;
                    for (int nxt : nxts) {
                        if (nxt != fa) {
                            best = Math.max(best, 1 + dist.get(nxt));
                        }
                    }
                    dist.put(node, best);
                }
            }
            return dist.get(start);
        }
    }
    
    
  • /**
     * Definition for a binary tree node.
     * struct TreeNode {
     *     int val;
     *     TreeNode *left;
     *     TreeNode *right;
     *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
     *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
     *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
     * };
     */
    class Solution {
    public:
        unordered_map<int, vector<int>> g;
    
        int amountOfTime(TreeNode* root, int start) {
            dfs(root);
            queue<int> q{ {start} };
            unordered_set<int> vis;
            int ans = -1;
            while (q.size()) {
                ++ans;
                for (int n = q.size(); n; --n) {
                    int i = q.front();
                    q.pop();
                    vis.insert(i);
                    for (int j : g[i]) {
                        if (!vis.count(j)) {
                            q.push(j);
                        }
                    }
                }
            }
            return ans;
        }
    
        void dfs(TreeNode* root) {
            if (!root) return;
            if (root->left) {
                g[root->val].push_back(root->left->val);
                g[root->left->val].push_back(root->val);
            }
            if (root->right) {
                g[root->val].push_back(root->right->val);
                g[root->right->val].push_back(root->val);
            }
            dfs(root->left);
            dfs(root->right);
        }
    };
    
    
    // Solution 2
    /**
     * Definition for a binary tree node.
     * struct TreeNode {
     *     int val;
     *     TreeNode *left;
     *     TreeNode *right;
     *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
     *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
     *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
     * };
     */
    class Solution {
    public:
        int amountOfTime(TreeNode* root, int start) {
            unordered_map<int, vector<int>> g;
            vector<pair<TreeNode*, TreeNode*>> stk{{root, nullptr}};
            while (!stk.empty()) {
                auto [node, fa] = stk.back();
                stk.pop_back();
                if (!node) {
                    continue;
                }
                if (fa) {
                    g[node->val].push_back(fa->val);
                    g[fa->val].push_back(node->val);
                }
                stk.emplace_back(node->right, node);
                stk.emplace_back(node->left, node);
            }
            unordered_map<int, int> dist;
            vector<tuple<int, int, int>> walk{{start, -1, 0}};
            while (!walk.empty()) {
                auto [node, fa, state] = walk.back();
                walk.pop_back();
                auto& nxts = g[node];
                if (state == 0) {
                    walk.emplace_back(node, fa, 1);
                    for (int i = (int) nxts.size() - 1; i >= 0; --i) {
                        int nxt = nxts[i];
                        if (nxt != fa) {
                            walk.emplace_back(nxt, node, 0);
                        }
                    }
                } else {
                    int best = 0;
                    for (int nxt : nxts) {
                        if (nxt != fa) {
                            best = max(best, 1 + dist[nxt]);
                        }
                    }
                    dist[node] = best;
                }
            }
            return dist[start];
        }
    };
    
    
  • # Definition for a binary tree node.
    # class TreeNode:
    #     def __init__(self, val=0, left=None, right=None):
    #         self.val = val
    #         self.left = left
    #         self.right = right
    class Solution:
        def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
            def dfs(root):
                if root is None:
                    return
                if root.left:
                    g[root.val].append(root.left.val)
                    g[root.left.val].append(root.val)
                if root.right:
                    g[root.val].append(root.right.val)
                    g[root.right.val].append(root.val)
                dfs(root.left)
                dfs(root.right)
    
            g = defaultdict(list)
            dfs(root)
            vis = set()
            q = deque([start])
            ans = -1
            while q:
                ans += 1
                for _ in range(len(q)):
                    i = q.popleft()
                    vis.add(i)
                    for j in g[i]:
                        if j not in vis:
                            q.append(j)
            return ans
    
    
    
    # Solution 2
    # Definition for a binary tree node.
    # class TreeNode:
    #     def __init__(self, val=0, left=None, right=None):
    #         self.val = val
    #         self.left = left
    #         self.right = right
    class Solution:
        def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
            g = defaultdict(list)
            stk = [(root, None)]
            while stk:
                node, fa = stk.pop()
                if node is None:
                    continue
                if fa:
                    g[node.val].append(fa.val)
                    g[fa.val].append(node.val)
                stk.append((node.right, node))
                stk.append((node.left, node))
    
            dist = {}
            walk = [(start, -1, 0)]
            while walk:
                node, fa, state = walk.pop()
                nxts = g[node]
                if state == 0:
                    walk.append((node, fa, 1))
                    for nxt in reversed(nxts):
                        if nxt != fa:
                            walk.append((nxt, node, 0))
                    continue
                best = 0
                for nxt in nxts:
                    if nxt != fa:
                        best = max(best, 1 + dist[nxt])
                dist[node] = best
            return dist[start]
    
    
  • /**
     * Definition for a binary tree node.
     * type TreeNode struct {
     *     Val int
     *     Left *TreeNode
     *     Right *TreeNode
     * }
     */
    func amountOfTime(root *TreeNode, start int) int {
    	g := map[int][]int{}
    	var dfs func(*TreeNode)
    	dfs = func(root *TreeNode) {
    		if root == nil {
    			return
    		}
    		if root.Left != nil {
    			g[root.Val] = append(g[root.Val], root.Left.Val)
    			g[root.Left.Val] = append(g[root.Left.Val], root.Val)
    		}
    		if root.Right != nil {
    			g[root.Val] = append(g[root.Val], root.Right.Val)
    			g[root.Right.Val] = append(g[root.Right.Val], root.Val)
    		}
    		dfs(root.Left)
    		dfs(root.Right)
    	}
    
    	dfs(root)
    	q := []int{start}
    	ans := -1
    	vis := map[int]bool{}
    	for len(q) > 0 {
    		ans++
    		for n := len(q); n > 0; n-- {
    			i := q[0]
    			q = q[1:]
    			vis[i] = true
    			for _, j := range g[i] {
    				if !vis[j] {
    					q = append(q, j)
    				}
    			}
    		}
    	}
    	return ans
    }
    
    
    // Solution 2
    /**
     * Definition for a binary tree node.
     * type TreeNode struct {
     *     Val int
     *     Left *TreeNode
     *     Right *TreeNode
     * }
     */
    func amountOfTime(root *TreeNode, start int) int {
    	g := map[int][]int{}
    	type frame struct {
    		node, fa *TreeNode
    	}
    	stk := []frame{{root, nil}}
    	for len(stk) > 0 {
    		cur := stk[len(stk)-1]
    		stk = stk[:len(stk)-1]
    		node, fa := cur.node, cur.fa
    		if node == nil {
    			continue
    		}
    		if fa != nil {
    			g[node.Val] = append(g[node.Val], fa.Val)
    			g[fa.Val] = append(g[fa.Val], node.Val)
    		}
    		stk = append(stk, frame{node.Right, node}, frame{node.Left, node})
    	}
    	dist := map[int]int{}
    	type step struct {
    		node, fa, state int
    	}
    	walk := []step{{start, -1, 0}}
    	for len(walk) > 0 {
    		cur := walk[len(walk)-1]
    		walk = walk[:len(walk)-1]
    		node, fa, state := cur.node, cur.fa, cur.state
    		nxts := g[node]
    		if state == 0 {
    			walk = append(walk, step{node, fa, 1})
    			for i := len(nxts) - 1; i >= 0; i-- {
    				if nxt := nxts[i]; nxt != fa {
    					walk = append(walk, step{nxt, node, 0})
    				}
    			}
    			continue
    		}
    		best := 0
    		for _, nxt := range nxts {
    			if nxt != fa {
    				best = max(best, 1+dist[nxt])
    			}
    		}
    		dist[node] = best
    	}
    	return dist[start]
    }
    
    
  • /**
     * Definition for a binary tree node.
     * class TreeNode {
     *     val: number
     *     left: TreeNode | null
     *     right: TreeNode | null
     *     constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) {
     *         this.val = (val===undefined ? 0 : val)
     *         this.left = (left===undefined ? null : left)
     *         this.right = (right===undefined ? null : right)
     *     }
     * }
     */
    
    function amountOfTime(root: TreeNode | null, start: number): number {
        const map = new Map<number, number[]>();
        const create = ({ val, left, right }: TreeNode) => {
            if (left != null) {
                map.set(val, [...(map.get(val) ?? []), left.val]);
                map.set(left.val, [...(map.get(left.val) ?? []), val]);
                create(left);
            }
            if (right != null) {
                map.set(val, [...(map.get(val) ?? []), right.val]);
                map.set(right.val, [...(map.get(right.val) ?? []), val]);
                create(right);
            }
        };
        create(root);
        const dfs = (st: number, fa: number) => {
            let res = 0;
            for (const v of map.get(st) ?? []) {
                if (v !== fa) {
                    res = Math.max(res, dfs(v, st) + 1);
                }
            }
            return res;
        };
        return dfs(start, -1);
    }
    
    
    
    // Solution 2
    /**
     * Definition for a binary tree node.
     * class TreeNode {
     *     val: number
     *     left: TreeNode | null
     *     right: TreeNode | null
     *     constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) {
     *         this.val = (val===undefined ? 0 : val)
     *         this.left = (left===undefined ? null : left)
     *         this.right = (right===undefined ? null : right)
     *     }
     * }
     */
    
    function amountOfTime(root: TreeNode | null, start: number): number {
        const g: Map<number, number[]> = new Map();
        const stk: [TreeNode | null, TreeNode | null][] = [[root, null]];
        while (stk.length) {
            const [node, fa] = stk.pop()!;
            if (!node) {
                continue;
            }
            if (fa) {
                if (!g.has(node.val)) {
                    g.set(node.val, []);
                }
                g.get(node.val)!.push(fa.val);
                if (!g.has(fa.val)) {
                    g.set(fa.val, []);
                }
                g.get(fa.val)!.push(node.val);
            }
            stk.push([node.right, node]);
            stk.push([node.left, node]);
        }
        const dist = new Map<number, number>();
        const walk: [number, number, number][] = [[start, -1, 0]];
        while (walk.length) {
            const [node, fa, state] = walk.pop()!;
            const nxts = g.get(node) || [];
            if (state === 0) {
                walk.push([node, fa, 1]);
                for (let i = nxts.length - 1; i >= 0; --i) {
                    const nxt = nxts[i];
                    if (nxt !== fa) {
                        walk.push([nxt, node, 0]);
                    }
                }
                continue;
            }
            let best = 0;
            for (const nxt of nxts) {
                if (nxt !== fa) {
                    best = Math.max(best, 1 + dist.get(nxt)!);
                }
            }
            dist.set(node, best);
        }
        return dist.get(start)!;
    }
    
    

All Problems

All Solutions