Print Nodes at Distance K From a Given Node

Hard
0/120
Average time to solve is 20m
profile
Contributed by
70 upvotes
Asked in companies
Paytm (One97 Communications Limited)AppleAmazon

Problem statement

You are given an arbitrary binary tree, a node of the tree, and an integer 'K'. You need to find all such nodes which have a distance K from the given node and return the list of these nodes.


Distance between two nodes in a binary tree is defined as the number of connections/edges in the path between the two nodes.


Note:

1. A binary tree is a tree in which each node has at most two children. 
2. The given tree will be non-empty.
3. The given tree can have multiple nodes with the same value.
4. If there are no nodes in the tree which are at distance = K from the given node, return an empty list.
5. You can return the list of values of valid nodes in any order. For example if the valid nodes have values 1,2,3, then you can return {1,2,3} or {3,1,2} etc.
Example :

Sample Output 2 explanation

Consider this tree above. The target node is 5 and K = 3. The nodes at distance 1 from node 5 are {2}, nodes at distance 2 from node 5 are {1, 4} and nodes at distance 3 from node 5 are {6, 3}.
Detailed explanation ( Input/output format, Notes, Images )
Input Format:
The first line will contain the values of the nodes of the tree in the level order form ( -1 for NULL node). Refer to the example below for further explanation.

The second line contains the value of the target node.

The third and the last line contains the integer K denoting the distance at which nodes are to be found. 

Example:

Consider the binary tree:

Input Format Fig.

The input for the tree depicted in the above image would be :

3
5 1
6 2 0 8
-1 -1 7 4 -1 -1 -1 -1
-1 -1 -1 -1

Explanation :
Level 1 :
The root node of the tree is 3

Level 2 :
Left child of 3 = 5
Right child of 3 = 1

Level 3 :
Left child of 5 = 6
Right child of 5 = 2
Left child of 1 = 0
Right child of 1 = 8

Level 4 :
Left child of 6 = null (-1)
Right child of 6 = null(-1)
Left child of 2 = 7
Right child of 2 = 4
Left child of 0 = null (-1)
Right child of 0 = null (-1)
Left child of 8 = null (-1)
Right child of 8 = null (-1)

Level 5 :
Left child of 7 = null (-1)
Right child of 7 = null (-1)
Left child of 4 = null (-1)
Right child of 4 = null (-1)

The first not-null node (of the previous level) is treated as the parent of the first two nodes of the current level. The second not-null node (of the previous level) is treated as the parent node for the next two nodes of the current level and so on.

The input ends when all nodes at the last level are null (-1).

Note :

The above format was just to provide clarity on how the input is formed for a given tree. 
The sequence will be put together in a single line separated by a single space. Hence, for the above-depicted tree, the input will be given as:

3 5 1 6 2 0 8 -1 -1 7 4 -1 -1 -1 -1 -1 -1 -1 -1
Output Format :
Print the values of all nodes at distance = K, from the given target node.

Note:

You do not need to print anything, it has already been taken care of. Just implement the given function.
Sample Input 1 :
3 5 1 6 2 0 8 -1 -1 7 4 -1 -1 -1 -1 -1 -1 -1 -1
5
2
Sample Output 1 :
7 4 1
Explanation For The Sample Output 1:

Sample Input 1 explanation

Target Node is 5. Nodes at distance 1 from 5 are {6, 2, 3} and nodes at distance 2 are {7, 4, 1}.
Sample Input 2:
7 -1 14 9 -1 13 -1 20 9 -1 8 -1 -1 2 -1 -1 16 -1 -1 
2
6
Sample Output 2:
7
Constraints:
1 <= N <= 3000
0 <= K <= 3000
0 <= nodeValue <= 3000

Where nodeValue donates the value of the node.

Time Limit: 1 sec
Hint

Store parents of each node in the tree.

Approaches (2)
DFS
  • Create a map to store the parent of each node in the tree, Traverse the tree recursively (via depth-first search), at each step if the current node is not NULL. Store its parent in the map, then traverse the left and right subtree.
  • Now assume that the given node is the root of the tree. In such a case, we can simply run a breadth-first search from the root, and track the current level of the tree. When the level = ‘K’, then the current nodes in the queue will be our answer.
  • But the problem is if the target node is not the root, then we can’t travel to every node with only left and right pointers. So here the stored parents will help us.
  • Observe that in a binary tree a node can be connected to maximum 3 other nodes ie. left child, right child, and the parent. We have left and right pointers, and we have also stored the parent node.
  • Now simply start BFS from the target node, and for each node which is at front of the queue, push its left child, right child and parent node(provided they are not null).
  • When the level of BFS reaches ‘K’, break the BFS and store all the values of nodes in the queue to an array and return it.
Time Complexity

O(N), where ‘N’ is the number of nodes in the given tree.

 

As we are checking for every node once, hence the time complexity will be O(N).

Space Complexity

O(N), where ‘N’ is the number of nodes in the given tree.

 

As we use additional space in BFS for the queue, and in the recursion call stack for DFS. Hence the overall space complexity will be O(N).

Video Solution
Unlock at level 3
(75% EXP penalty)
Code Solution
(100% EXP penalty)
Print Nodes at Distance K From a Given Node
All tags
Sort by
Search icon

Interview problems

Easy || C++ implementation || Striver's Approach

Do Upvote if you found it helpful!

/*************************************************************

    Following is the Binary Tree node structure for reference:

    class BinaryTreeNode{
	public :
        T data;
        BinaryTreeNode<T> *left;
        BinaryTreeNode<T> *right;

        BinaryTreeNode(T data) {
            this -> data = data;
            left = NULL;
            right = NULL;
        }
    };

*************************************************************/
#include<bits/stdc++.h>
void mark_parents(BinaryTreeNode<int>*root,unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*>& parents,BinaryTreeNode<int>*target){
    queue<BinaryTreeNode<int>*>q;
    q.push(root);
    while(!q.empty()){
        BinaryTreeNode<int>* current = q.front();
        q.pop();
        if(current->left){
            parents[current->left]=current;
            q.push(current->left);
        }
        if(current->right){
            parents[current->right]=current;
            q.push(current->right);
        }
    }

}
vector<BinaryTreeNode<int>*> printNodesAtDistanceK(BinaryTreeNode<int>* root, BinaryTreeNode<int>* target, int K) {
    // Write your code here.
    vector<BinaryTreeNode<int>*>ans;
    unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*>parents;
    mark_parents(root,parents,target);
    unordered_map<BinaryTreeNode<int>*,bool>visited;
    queue<BinaryTreeNode<int>*>q;
    q.push(target);
    visited[target] = true;
    int curLevel = 0;
    while(!q.empty()){
        int size = q.size();
        if(curLevel++ == K)break;
        for(int i = 0; i<size ; i++){
            BinaryTreeNode<int>* current = q.front();
            q.pop();
            if(current->left&&!visited[current->left]){
                q.push(current->left);
                visited[current->left]=true;
            }
             if(current->right&&!visited[current->right]){
                q.push(current->right);
                visited[current->right]=true;
            }
             if(parents[current]&&!visited[parents[current]]){
                q.push(parents[current]);
                visited[parents[current]]=true;
            }
        }
    }
    while(!q.empty()){
        BinaryTreeNode<int>* current = q.front();
        q.pop();
        ans.push_back(current);
    }
    return ans;

}
61 views
0 replies
1 upvote

Interview problems

Python Easy Solution

def printNodesAtDistanceK(root, target, K):

    # Write your code here.

    def dfs(node, parent=None):

        if node:

            node.parent = parent

            dfs(node.left, node)

            dfs(node.right, node)

    

    dfs(root)

    

    # Breadth First Search from target to find nodes at distance K

    queue = [(target, 0)]

    visited = {target}

    result = []

    while queue:

        node, distance = queue.pop(0)

        if distance == K:

            result.append(node)

        elif distance > K:

            break

        

        neighbors = [node.left, node.right, node.parent]

        for neighbor in neighbors:

            if neighbor and neighbor not in visited:

                queue.append((neighbor, distance + 1))

                visited.add(neighbor)

    

    return result

 

54 views
0 replies
0 upvotes

Interview problems

C++ || Using map and bfs

void createMapping(BinaryTreeNode<int>* root, map<BinaryTreeNode<int>*, BinaryTreeNode<int>*> &nodeToParent) {

    

    if (root == nullptr)

        return;

 

    queue<BinaryTreeNode<int>*> q;

    q.push(root);

    nodeToParent[root] = nullptr;

 

    while (!q.empty()) {

 

        BinaryTreeNode<int>* front = q.front();

        q.pop();

 

        if (front->left) {

            nodeToParent[front->left] = front;

            q.push(front->left);

        }

 

        if (front->right) {

            nodeToParent[front->right] = front;

            q.push(front->right);

        }

    }

}

 

vector<BinaryTreeNode<int>*> printNodesAtDistanceK(BinaryTreeNode<int>* root, BinaryTreeNode<int>* target, int K) {

    

    vector<BinaryTreeNode<int>*> ans;

 

    if (root == NULL || target == NULL)

        return ans;

 

    // Step 1: Creating a node to parent mapping

    map<BinaryTreeNode<int>*, BinaryTreeNode<int>*> nodeToParent;

    createMapping(root, nodeToParent);

 

    // Step 2: Do BFS and find the level K and insert nodes

    map<BinaryTreeNode<int>*, bool> visited;

    queue<BinaryTreeNode<int>*> q;

    q.push(target);

    visited[target] = 1;

 

    while(!q.empty() && K >= 0){

 

        int n = q.size();

 

        while(n--){

 

            BinaryTreeNode<int> *front = q.front();

            q.pop();

 

            if(K == 0) ans.push_back(front);

 

            if(front->left && !visited[front->left]){

 

                q.push(front->left);

                visited[front->left] = 1;

 

            }

 

            if(front->right && !visited[front->right]){

 

                q.push(front->right);

                visited[front->right] = 1;

 

            }

 

            if(nodeToParent[front] && !visited[nodeToParent[front]]){

 

                q.push(nodeToParent[front]);

                visited[nodeToParent[front]] = 1;

 

            }

 

        }

 

        // one level covered

        K--;

 

    }

 

    return ans;

}

 

356 views
0 replies
4 upvotes

Interview problems

fast and best solution...

void parent(BinaryTreeNode<int>* root,unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*>& mp){

    queue<BinaryTreeNode<int>*> q;

    q.push(root);

    while(!q.empty()){

        BinaryTreeNode<int>* 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);

        }

    }

}

vector<BinaryTreeNode<int>*> printNodesAtDistanceK(BinaryTreeNode<int>* root, BinaryTreeNode<int>* target, int K) {

    unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*> mp;

    parent(root,mp);

    queue<BinaryTreeNode<int>* > q;

 

    unordered_map<BinaryTreeNode<int>* ,bool> check;

 

    int cnt = 0;

    q.push(target);

    check[target] = true;

    while(!q.empty()){

        int size = q.size();

        if((cnt++) == K)break;

        for(int i =0;i < size; i++){

            BinaryTreeNode<int>* node = q.front();

            q.pop();

            if(node->left && !check[node->left]){

                q.push(node->left);

                check[node->left] = true;

            }

            if(node->right && !check[node->right]){

                q.push(node->right);

                check[node->right] = true;

            }

            if(mp[node] && !check[mp[node]]){

                q.push(mp[node]);

                check[mp[node]] = true;

            }

        }

    }

    vector<BinaryTreeNode<int>*> vec;

   while(!q.empty()){

       vec.push_back(q.front());

       q.pop();

   }

 

return vec;

}

204 views
0 replies
0 upvotes

Interview problems

Java BFS Solution

import java.util.*;
public class Solution {
    public static void addParent(BinaryTreeNode<Integer> root, BinaryTreeNode<Integer> p, Map<BinaryTreeNode<Integer>, BinaryTreeNode<Integer>> hm){
      if(root==null) return;
      hm.put(root,p);
      addParent(root.left,root,hm);
      addParent(root.right,root,hm);
    }
    public static List<BinaryTreeNode<Integer>> printNodesAtDistanceK(BinaryTreeNode<Integer> root, BinaryTreeNode<Integer> target, int K) {
        // Write your code here.
        List<BinaryTreeNode<Integer>> al=new ArrayList<>();
        Queue<BinaryTreeNode<Integer>> q=new LinkedList<>();
        Map<BinaryTreeNode<Integer>, BinaryTreeNode<Integer>> parent=new HashMap<>();
        Set<BinaryTreeNode<Integer>> vis=new HashSet<>();
        addParent(root,null,parent);
        q.add(target);
        vis.add(target);
        while(!q.isEmpty()){
         int n=q.size();
         for(int i=0;i<n;i++){
          BinaryTreeNode<Integer> curr=q.poll();
          if(K==0) al.add(curr);
          if(curr.left!=null && !vis.contains(curr.left)) {
              q.add(curr.left);
              vis.add(curr.left);
          }
          if(curr.right!=null && !vis.contains(curr.right)) {
              q.add(curr.right);
              vis.add(curr.right);
          }
          if(parent.get(curr)!=null && !vis.contains(parent.get(curr))){
              q.add(parent.get(curr));
              vis.add(parent.get(curr));
          }
          }
          K--;
          if(K<0) break;
        }
        return al;
    }
}
153 views
0 replies
0 upvotes

Interview problems

C++ EASY TO IMPLEMENT SOL || LINEAR TIME🔥🏆

#include<bits/stdc++.h>

void x(BinaryTreeNode<int>*root,BinaryTreeNode<int>*target,int k,vector<BinaryTreeNode<int>*>&ans,unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*> mp,unordered_map<BinaryTreeNode<int>*,bool> &mpx){

   queue<pair<BinaryTreeNode<int>*,int>> q;

   q.push({target,k});

 

   while(!q.empty()){

       BinaryTreeNode<int>*temp=q.front().first;

       int l=q.front().second;

       q.pop();

 

       mpx[temp]=1;

    if(l==0){

        ans.push_back(temp);

    }

 

       if(temp->left!=NULL && mpx[temp->left]==0){

           q.push({temp->left,l-1});

       }

        if(temp->right!=NULL && mpx[temp->right]==0){

           q.push({temp->right,l-1});

       }

        if(mp[temp]!=NULL && mpx[mp[temp]]==0){

           q.push({mp[temp],l-1});

       }

   }

 

}

 

vector<BinaryTreeNode<int>*> printNodesAtDistanceK(BinaryTreeNode<int>* root, BinaryTreeNode<int>* target, int K) {

   if(root==NULL)return {};

   if(target==NULL)return {};

   if(K==0)return {target};

   unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*> mp;

   unordered_map<BinaryTreeNode<int>*,bool> mpx;

   queue<BinaryTreeNode<int>*> q;

 

   q.push(root);

vector<BinaryTreeNode<int>*>ans;

   while(!q.empty()){

       BinaryTreeNode<int>* temp=q.front();

       q.pop();

       if(temp->left){

           q.push(temp->left);

           mp[temp->left]=temp;

       }

        if(temp->right){

           q.push(temp->right);

           mp[temp->right]=temp;

       }

   }

 

   x(target,target,K,ans,mp,mpx);

   return ans;

 

}

213 views
0 replies
0 upvotes

Interview problems

C++ Solution || (print Node at dist K )||stack

//traverse node to find target node and store path in stack
void trav(BinaryTreeNode<int>* root, BinaryTreeNode<int>* target,bool &find,stack<BinaryTreeNode<int>*> &s){
    if(root==NULL)
        return ;
    if(root==target){
        find=true;
        return;
    }
    s.push(root->left);
    trav(root->left,target,find,s);
    if(find){
        return;
    }
    s.pop();
    s.push(root->right);
    trav(root->right,target,find,s);
    if(find){
        return;
    }
    s.pop();
}
// finding node at distance k from a node 
void fkd(BinaryTreeNode<int>* root,BinaryTreeNode<int>* past,int K,vector<BinaryTreeNode<int>*> &ans){
    if(root==NULL||root==past)
        return;
    if(K==0){
        ans.push_back(root);
        return;
    }
    
    fkd(root->left,past,K-1,ans);
    fkd(root->right,past,K-1,ans);
    
}
vector<BinaryTreeNode<int>*> printNodesAtDistanceK(BinaryTreeNode<int>* root, BinaryTreeNode<int>* target, int K) {
    // Write your code here.
    vector<BinaryTreeNode<int>*> ans;
    stack<BinaryTreeNode<int>*> s;
    s.push(root);
    bool find=false;
    trav(root,target,find,s);
    BinaryTreeNode<int>* past=NULL;
    while(!s.empty()){
        BinaryTreeNode<int>* front=s.top();
        s.pop();
        fkd(front,past,K,ans);
        past=front;
        K--;
    }
    return ans ;

}
147 views
0 replies
1 upvote

Interview problems

JAVA Solution (Print Nodes at Distance K From a Given Node)

/*************************************************************

	Following is the Binary Tree node structure for reference:
 
  	class BinaryTreeNode<T> {
  
  		T data; 
		BinaryTreeNode<T> left; 
		BinaryTreeNode<T> right;
  
  		BinaryTreeNode(T data) { 
			this.data = data; 
			left = null; right = null; 
		} 
	}

*************************************************************/

import java.util.*;
class found{
    boolean path;
    int count;
    public found(boolean a,int b){
        path=a;
        count=b;
    }
}
public class Solution {	
    static HashMap<BinaryTreeNode<Integer>,Integer> map=new HashMap<>();
    static BinaryTreeNode<Integer> tofind;
    static int dis;
    static ArrayList<BinaryTreeNode<Integer>> ret=new ArrayList<>();
    public static found udfs(BinaryTreeNode<Integer> root){
        if(root==null){
            return new found(false,0);
        }
        if(root==tofind){
            addlist(root,0);
            map.put(root,0);
            return new found(true,0);
        }
        found f1=udfs(root.right);
        if(f1.path==true){
            int temp=f1.count;
            if(temp<=dis){
                addlist(root,temp+1);
                map.put(root,1);
                return new found(true,temp+1);
            }
        }
        found f2=udfs(root.left);
        if(f2.path==true){
            int temp=f2.count;
            if(temp<=dis){
                addlist(root,temp+1);
                map.put(root,1);
                return new found(true,temp+1);
            }
        }
       return new found(false,0); 
    }
    public static void addlist(BinaryTreeNode<Integer> root,int n){
        if(root==null){
            return;
        }
        if(n==dis){
            ret.add(root);
            return;
        }
        if(!map.containsKey(root.left)){
            addlist(root.left,n+1); 
        }
       if(!map.containsKey(root.right)){
            addlist(root.right,n+1); 
        }
    }
	public static ArrayList<BinaryTreeNode<Integer>> printNodesAtDistanceK(BinaryTreeNode<Integer> root,
   		BinaryTreeNode<Integer> target, int K) {
        tofind=target;
        dis=K;
        ret.clear();
        map.clear();
		udfs(root);
        return ret;
	}
}
71 views
0 replies
0 upvotes

Interview problems

Using DFS || Map || queue

void findparent(BinaryTreeNode<int>* root,unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*>& ml){

    if(root==NULL) return;

    if(root->left) ml[root->left]=root;

    if(root->right) ml[root->right]=root;

    findparent(root->left,ml);

    findparent(root->right,ml);

}

vector<BinaryTreeNode<int>*> printNodesAtDistanceK(BinaryTreeNode<int>* root, BinaryTreeNode<int>* target, int K) {

    unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*> ml;

    findparent(root,ml);

    unordered_map<BinaryTreeNode<int>*,bool>visited;

    queue<BinaryTreeNode<int>*>q;

    q.push(target);

    visited[target]=true;

    int curr=0;

    while(!q.empty()){

        int size=q.size();

        if(curr++==K) break;

        for(int i=0;i<size;i++){

            BinaryTreeNode<int>* current=q.front(); q.pop();

            if(current->left && !visited[current->left]){

                q.push(current->left);

                visited[current->left]=true;

            }

            if(current->right && !visited[current->right]){

                q.push(current->right);

                visited[current->right]=true;

            }

            if(ml[current] && !visited[ml[current]]){

                q.push(ml[current]);

                visited[ml[current]]=true;

            }

        }

    }

    vector<BinaryTreeNode<int>*>ans;

    while(!q.empty()){

        ans.push_back(q.front());

        q.pop();

    }

    return ans;

}

134 views
0 replies
1 upvote

Interview problems

finding node parent from dfs traversal

void findparent(BinaryTreeNode<int>* root,unordered_map<BinaryTreeNode<int>*,BinaryTreeNode<int>*>& ml)

{

    if(root==NULL) return;

    if(root->left) ml[root->left]=root;

    if(root->right) ml[root->right]=root;

    findparent(root->left,ml);

    findparent(root->right,ml);

}

 

68 views
0 replies
0 upvotes
Full screen
Console