863. All Nodes Distance K in Binary Tree

#binary-trees

Problem

All Nodes Distance K in Binary Tree - LeetCodeLeetCode

Intuition

Time Complexity

O(V + E) -> BFS for marking parents
O(V + E) -> BFS for traversing from target
Worst case is that taregt is root node and the "k" is max depth of the tree

Total Time Complexity = O(V + E) + O(V + E) = 2 * O(V + E) = O(V + E)

Space Complexity

O(V + E), to store the parent child mapping

Solution

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    private Map<TreeNode, TreeNode> markParents(TreeNode node) {
        // Standard BFS traversal
        Deque<TreeNode> q = new ArrayDeque<>();
        q.offer(node);

        // Create a parents map
        Map<TreeNode, TreeNode> parentMap = new HashMap<>();

        while(!q.isEmpty()) {
            // The difference in this traversal is that, we do not go level by
            // level exactly. 
            TreeNode current = q.poll();
            
            // At every node, we just check if the left subtree
            // and right subtree is null or not
            if(current.left!=null) {
                // If it is not null, then in the parent map, we just add the
                // "current left node" parent as the "current" node
                parentMap.put(current.left, current);
                q.offer(current.left);
            }

            if(current.right != null) {
                // If it is not null, then in the parent map, we just add the
                // "current right node" parent as the "current" node
                parentMap.put(current.right, current);
                q.offer(current.right);
            }
        }

        return parentMap;
    }

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        // Time Complexity:
        // O(V + E) -> BFS for marking parents
        // O(V + E) -> BFS for traversing from target
        // Worst case is that taregt is root node and the "k" is max depth of the tree
        // Total Time Complexity = O(V + E) + O(V + E) = 2 * O(V + E) = O(V + E)

        // Space Complexity: O(V + E), to store the parent child mapping

        // Create a map which has information about every nodes parents
        Map<TreeNode, TreeNode> parents = markParents(root);

        // Here the traversal is not only one way i.e. bottom only, but also to top
        // So we include parent nodes whiles traversing and to make sure the 
        // node is not double processed, we create a "visited" set which keeps
        // a track of node which have been already processed
        Set<TreeNode> visited = new HashSet<>();

        // BFS traversal
        Deque<TreeNode> q = new ArrayDeque<>();
        // We directly add the target node to the queue. That will be the starting
        // point, this makes sure that we will traverse only "k" units distance from
        // the tree
        q.offer(target);
        visited.add(target);

        int currentLevel = 0;

        while(!q.isEmpty()) {
            int n = q.size();
            
            // When the current level equals to the "k", then we 
            // do not need to process anything else, so break
            if(currentLevel == k) break;

            for(int i=0;i<n;i++) {
                TreeNode current = q.poll();

                // Get the left, right and parent node
                TreeNode left = current.left;
                TreeNode right = current.right;
                TreeNode parent = parents.get(current);

                // Check if they are null or not and if they are visited or not
                if(left != null && !visited.contains(left)) {
                    // If they pass the above two conditions, then
                    // add to queue for processing
                    q.offer(left);
                    // Add to visited
                    visited.add(left);
                }

                // Do same as what was done for left subtree
                if(right != null && !visited.contains(right)) {
                    q.offer(right);
                    visited.add(right);
                }

                // Do same as what was done for left subtree
                if(parent!=null && !visited.contains(parent)) {
                    q.offer(parent);
                    visited.add(parent);
                }
            }
            // Increment the level after each level
            currentLevel++;
        }

        // Create a result container
        List<Integer> result = new ArrayList<>();

        // Once we reach the specified distance by using the property 
        // currentLevel == k
        // Then the "q" contains only the nodes which are at distance "k"
        // traverse through the queue and add the values to the result
        for(TreeNode curr: q) result.add(curr.val);
               
        return result;
    }
}