DAY 60 - All Nodes Distance K in Binary Tree

Welcome to DAY 60. Today we will diverge once again into the Data Structure Trees and look at some more questions. I suggest reading the previous articles that I have written on Trees before reading this.

Question: All Nodes Distance K in Binary Tree

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer in any order.

Example 1:

• Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
• Output: [7,4,1]
• Explanation: The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.

Intuition:

• Well the question is pretty easy. Let's take a look at it.
• Imagine our starting point is the root instead of target for once.
• If we need to find all the nodes at K distance, we can simply do bfs till K distance as it travels all nearest nodes at a unit distance and we could store them. That would be our answer.
• The catch here is that we are given any node of the tree.
• We can apply BFS traversal there but the issue which comes next is how do we travel up.
• Take 5 for example, we can travel to it's left and right nodes 6 and 2 but we can't go to 1.
• If we can devise a way to go to parent node back from the child node, we have all the things we need.
• Then we can just to bfs traversal here and find the answer.

Algorithm and Code:

Algorithm:

• The parentMap function creates a parent map (mp) for each node in the binary tree using a breadth-first search (BFS) traversal. The parent map stores the parent of each node in the tree, and it will be later used to traverse back from the target node to find nodes at distance K from the target.
• The distanceK function takes the root of the binary tree, the target node, and the integer K as input and returns a vector of integers representing the nodes at distance K from the target.
• The function initializes an empty unordered_map mp to store the parent map and an unordered_map visited to keep track of visited nodes during BFS traversal.
• It starts a BFS traversal from the target node by pushing it into a queue q. The target node is marked as visited.
• The algorithm uses a variable cur_lvl to keep track of the current level in the BFS traversal. It dequeues nodes from the queue one level at a time, stopping the traversal once the desired distance K is reached.
• For each node in the current level, the algorithm enqueues its left child and right child if they exist and are not already visited. It also enqueues the parent of the current node (retrieved from the parent map mp) if the parent exists and is not already visited.
• After processing all nodes in the current level, the cur_lvl is incremented to move to the next level.
• Once the BFS traversal is complete, the algorithm stores the nodes at distance K from the target node in the result vector.
• Finally, the function returns the result vector containing the nodes at distance K from the target node in the binary tree.

Code:

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    // Function to create a parent map for each node in the binary tree
    void parentMap(TreeNode* root, unordered_map<TreeNode*, TreeNode*> &mp, TreeNode* target)
    {
        queue<TreeNode*> q;
        q.push(root);
        while(!q.empty()){
            TreeNode* node = q.front();
            q.pop();
            if(node->left) {
                mp[node->left] = node;
                q.push(node->left);
            }
            if(node->right) {
                mp[node->right] = node;
                q.push(node->right);
            }
        }
    }

    // Function to find nodes at distance K from the target node in the binary tree
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        unordered_map<TreeNode*, TreeNode*> mp;
        parentMap(root, mp, target);
        unordered_map<TreeNode*, bool> visited;
        queue<TreeNode*> q;
        q.push(target);
        visited[target] = true;
        int cur_lvl = 0;

        while(!q.empty()){
            int size = q.size();
            // If we have reached the desired distance K, stop exploring further
            if(cur_lvl++ == k) break;
            for(int i = 0; i < size; i++){
                TreeNode* current = q.front();
                q.pop();
                if(current->left && visited[current->left] == false){
                    q.push(current->left);
                    visited[current->left] = true;
                }
                if(current->right && visited[current->right] == false){
                    q.push(current->right);
                    visited[current->right] = true;
                }
                if(mp[current] && !visited[mp[current]]){
                    q.push(mp[current]);
                    visited[mp[current]] = true;
                }
            }
        }

        // Store the nodes at distance K from the target in the result vector
        vector<int> result;
        while(!q.empty()){
            TreeNode* current = q.front();
            q.pop();
            result.push_back(current->val);
        }

        return result;
    }
};

Complexity Analysis:

Time: O(N) + O(K)
Space: O(N) + O(K)