# Find all nodes at distance K from a given target node in a binary tree

• Onsite 2nd round Amazon SDE1 interview.

Given a binary tree, find all nodes at a distance K from a given target node.

• ``````package Amazon;

public class TreeNodesDistanceK {

private void printDownwardPaths(Node cur, int k) {
if(null == cur) {
return;
}

if(k == 0) {
System.out.println(cur.data);
}

printDownwardPaths(cur.left, k-1);
printDownwardPaths(cur.right, k-1);
}

private int printUpwardPaths(Node cur, Node target, int k) {
if(null == cur) {
return -1;
}

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

int dist_left = printUpwardPaths(cur.left, target, k);
int dist_right = printUpwardPaths(cur.right, target, k);

if(dist_left != -1) {
if(dist_left+1 == k) {
System.out.println(cur.data);
} else {
printDownwardPaths(cur.right, k-dist_left-2);
}

return dist_left+1;
}

if(dist_right != -1) {
if(dist_right+1 == k) {
System.out.println(cur.data);
} else {
printDownwardPaths(cur.left, k-dist_right-2);
}

return dist_right+1;
}

return -1;
}

private void printNodesAtDistanceK(Node root, Node target, int k) {
printDownwardPaths(target, k);
printUpwardPaths(root, target, k);
}

public static void main(String[] args) {
TreeNodesDistanceK tdk = new TreeNodesDistanceK();

/**
* 					    1
* 				      /          \
* 				    2	           3
*                              /      \        /     \
*                          4          5     6       7
*                       /     \   /     \
*                     8      9 10     11
*                                            \
*                                             12
*                                           /      \
*                                        13      14
* */

Node n1 = new Node(1);
Node n2 = new Node(2);
Node n3 = new Node(3);
Node n4 = new Node(4);
Node n5 = new Node(5);
Node n6 = new Node(6);
Node n7 = new Node(7);
Node n8 = new Node(8);
Node n9 = new Node(9);
Node n10 = new Node(10);
Node n11 = new Node(11);
Node n12 = new Node(12);
Node n13 = new Node(13);
Node n14 = new Node(14);

n1.left = n2;
n1.right = n3;

n2.left = n4;
n2.right = n5;

n3.left = n6;
n3.right = n7;

n4.left = n8;
n4.right = n9;

n5.left = n10;
n5.right = n11;

n11.left = n12;
n12.left = n13;
n12.right = n14;

tdk.printNodesAtDistanceK(n1, n5, 3);

}

}

``````

• ``````void printkdistanceNode(vector<int> v, int trg, int k){
int N=v.size();
if(N==0)return;
for(int i=0;i<N;++i){
if(v[i]==trg){
printkdistanceNodeDown(v,i,k);
int cur=i;
int parent=(cur-1)/2;
while(parent>=0){
k--;
if(i==(2*parent+1))printkdistanceNodeDown(v,2*parent+2,k-1);
if(i==(2*parent+2))printkdistanceNodeDown(v,2*parent+1,k-1);
cur=parent;
parent=(cur-1)/2;
}
return;
}
}
}

void printkdistanceNodeDown(vector<int> v, int root, int k){
if(k==0){
cout<<v[root]<<endl;
return;
}
int left=root*2+1;
int right=root*2+2;
if(left<v.size())printkdistanceNodeDown(v,left,k-1);
if(right<v.size())printkdistanceNodeDown(v,left,k-1);
}
``````

• @ramanpreetSinghKhinda
In method printUpwardPaths(), call to printUpwardPaths(cur.right, target, k) should not be made without checking the return value of printUpwardPaths(cur.left, target, k) first. The node is present only in either left or right subtree. If it is found in left-subtree, call should not be made to right sub-tree. Code as is will give correct result but does unnecessary processing.

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.