Print Tree nodes at distance K from random node

Problem:
Given pointer to root of the binary tree and pointer to random node in the tree. Find all nodes that are at distance K from the random node.

Solution:

Algorithm:
a. If root is the random node then print all nodes from the root that are at distance K. Printing of nodes can be done recursively.
b. Search the random node in left sub tree of root. If node is found in left sub-tree and let say distance of random node from root is X. If X <= K, then print all nodes in the right sub-tree of the root that are are at distance K - X from the root.
c. If node is not found in the left sub-tree of the root then search the random node in right sub-tree of the root. Let us say that distance of random node in right sub-tree is X from the root. If X <= K, then print all nodes in the left sub-tree of the root that are are at distance K - X from the root.
d. Move one step closer toward the random node. i.e. if random node exist in left sub-tree of the root then set root = root.left, otherwise set root = root.right. Repeat step (a) to step (c).

This algorithm described above will be O(n^2) but if we use recursion carefully then we can make same solution O(n). In same recursion keep searching the node by moving towards the node and also printing the nodes in relevant sub-tree.

Working code for O(n) solution.

 public class PrintNodesAtK {  
      public static void main(String[] args) {  
           Tree root = new Tree(1);  
           Tree p1 = new Tree(2);  
           Tree p2 = new Tree(3);  
           Tree p3 = new Tree(4);  
           Tree p4 = new Tree(5);  
           Tree p5 = new Tree(6);  
           Tree p6 = new Tree(7);  
           Tree p7 = new Tree(8);  
           Tree p8 = new Tree(9);  
           Tree p9 = new Tree(10);  
           Tree p10 = new Tree(11);  
           Tree p11 = new Tree(12);  
           Tree p12 = new Tree(13);  
           root.left = p1;  
           root.right = p2;  
           p1.left = p3;  
           p1.right = p4;  
           p2.left = p5;  
           p2.right = p6;  
           p3.left = p7;  
           p4.left = p8;  
           p4.right = p9;  
           p6.left = p10;  
           p6.right = p11;  
           p10.right = p12;  
           searchNode(root, p6, 5);  
      }  
      public static int searchNode(Tree root, Tree random, int k) {  
           if (k == 0 || root == null) {  
                return -1;  
           } else if (random == root) {  
                printNodeAtK(root, k);  
                return 0;  
           } else {  
                int p = searchNode(root.left, random, k);  
                if (p != -1 && p + 1 <= k) {  
                     // Print root if d = 0  
                     int d = k - p - 1;  
                     if (d == 0) {  
                          System.out.println(root.value);  
                     } else {  
                          printNodeAtK(root.right, d - 1);  
                     }  
                } else if (p == -1) {  
                     p = searchNode(root.right, random, k);  
                     if (p != -1 && p + 1 <= k) {  
                          // Print root if d = 0  
                          int d = k - p - 1;  
                          if (d == 0) {  
                               System.out.println(root.value);  
                          } else {  
                               printNodeAtK(root.left, d - 1);  
                          }  
                     }  
                }  
                return p == -1 ? -1 : p + 1;  
           }  
      }  
      public static void printNodeAtK(Tree root, int k) {  
           if(root == null) {  
                return;  
           } else if (k == 0) {  
                System.out.println(root.value);  
           } else {  
                printNodeAtK(root.left, k - 1);  
                printNodeAtK(root.right, k - 1);  
           }  
      }  
 }  
 class Tree {  
      int value;  
      Tree left;  
      Tree right;  
      public Tree(int value) {  
           this.value = value;  
           left = null;  
           right = null;  
      }  
 }