Amount of Time for Binary Tree to Be Infected (Burnt)

Problem 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. (link)

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.

Solution

Idea: The parent and child nodes catch fire at the same time.

When we are at any specific node, we need to know about three things:

• Parent node

• Left child node

• Right child node

We store the parents of each node in a map. This map provides the parent of any given node. We also keep track of the nodes that have already burned, so if we encounter a burned node again, we won't burn it again.

Note: We use BFS to build the parent map and to burn the tree.

Time - O(n)

Space - O(n)

class Solution {
    public int amountOfTime(TreeNode root, int start) {
        Map<TreeNode, TreeNode> parentMap = new HashMap<>();

        TreeNode startNode = buildParentMapAndFetchStartNode(root, parentMap, start);

        Set<TreeNode> burnt = new HashSet<>();
        Queue<TreeNode> queue = new ArrayDeque<>();
        queue.add(startNode);
        int timer = -1;

        while(!queue.isEmpty()){
            timer+=1;
            int size = queue.size();
            for(int i=0; i<size; i++){
                TreeNode node = queue.poll();
                TreeNode left = node.left;
                TreeNode right = node.right;
                TreeNode parent = parentMap.get(node);

                burnt.add(node);

                if(left!=null && !burnt.contains(left)){
                    queue.add(left);
                }
                if(right!=null && !burnt.contains(right)){
                    queue.add(right);
                }
                if(parent.val!=0 && !burnt.contains(parent)){
                    queue.add(parent);
                }
            }
        }
        return timer;
    }

    private TreeNode buildParentMapAndFetchStartNode(TreeNode root, Map<TreeNode, TreeNode> parentMap, int start){
        parentMap.put(root,new TreeNode(0,root,root));
        TreeNode startNode = null;

        Queue<TreeNode> queue = new ArrayDeque<>();
        queue.add(root);
        while (!queue.isEmpty()){
            TreeNode node = queue.poll();

            if(node.val == start){
                startNode = node;
            }

            if(node.left!=null){
                parentMap.put(node.left, node);
                queue.add(node.left);
            }

            if(node.right!=null){
                parentMap.put(node.right, node);
                queue.add(node.right);
            }
        }
        return startNode;
    }
}