Thursday, March 19, 2015

Print all nodes at distance k from a given node

Given a binary tree, a target node in the binary tree, and an integer value k, print all the nodes that are at distance k from the given target node. No parent pointers are available.
BinaryTree
Consider the tree shown in diagram

Input: target = pointer to node with data 8. 
       root = pointer to node with data 20.
       k = 2.
Output : 10 14 22

If target is 14 and k is 3, then output 
should be "4 20"



Reference: http://www.geeksforgeeks.org/print-nodes-distance-k-given-node-binary-tree/
Solution:
Two types of nodes to be considered.
1) Nodes in the subtree rooted with target node. For example if the target node is 8 and k is 2, then such nodes are 10 and 14. 
Traverse subtrees rooted with the target node and decrement k in recursive call. When the k becomes 0, print the node currently being traversed.
2) Other nodes, may be an ancestor of target, or a node in some other subtree. For target node 8 and k is 2, the node 22 comes in this category.
For the output nodes not lying in the subtree with the target node as the root, we must go through all ancestors. For every ancestor, we find its distance from target node, let the distance be d, now we go to other subtree (if target was found in left subtree, then we go to right subtree) of the ancestor and find all nodes at k-d distance from the ancestor.
Time Complexity is O(n).
public class TreeNode {
  int val;
  TreeNode left;
  TreeNode right;
  TreeNode(int x) { val = x; }
 }
 
 //Recursive function to print all the nodes at distance k in the tree (or subtree) rooted with given root.
 public static void printKthDown(TreeNode node, int k) {
     if(node == null || k < 0) {  //base case
         return;
     }
  
     if(k == 0) {  //If we reach a k distant node, print it
         System.out.println(node.val);
         return;
     }
   
     // Recur for left and right subtrees
     printKthDown(node.left, k - 1);
     printKthDown(node.right, k - 1);
 }
 
 // Prints all nodes at distance k from a given target node.
 // The k distant nodes may be upward or downward.  This function
 // Returns distance of root from target node, it returns -1 if target
 // node is not present in tree rooted with root.
 public static int printKthDistanceNode(TreeNode root, TreeNode target, int k) {
     if(root == null) {  // Base Case 1: If tree is empty, return -1
        return -1;
     }
    
     // If target is same as root.  Use the downward function
     // to print all nodes at distance k in subtree rooted with
     // target or root
     if(root == target) {
         printKthDown(root, k);
         return 0;
     }
  
     int leftDistance = printKthDistanceNode(root.left, target, k);  // Recur for left subtree
  
     if(leftDistance != -1) {
         // If root is at distance k from target, print root
         // Note that leftDistance is Distance of root's left child from target
         if(leftDistance + 1 == k) {
             System.out.println(root.val);
         }else {  //go to right subtree and print all k-leftDistance-2 distant nodes. Note that the right child is 2 edges away from left child
             printKthDown(root.right, k - leftDistance - 2);
         }
         return leftDistance + 1; // Add 1 to the distance and return value for parent calls
     }
  
     int rightDistance = printKthDistanceNode(root.right, target, k);
  
     if(rightDistance != -1) {
         if(rightDistance + 1 == k) {
             System.out.println(root.val);
         }else {
             printKthDown(root.left, k - rightDistance - 2);
         }
         return rightDistance + 1;
     }
  
     return -1;  // If target was neither present in left nor in right subtree
 }
完整代码:
package leetcode;

public class PrintAllNodeAtDistanceKFromANode {
 
 public class TreeNode {
  int val;
  TreeNode left;
  TreeNode right;
  TreeNode(int x) { val = x; }
 }
 
 //Recursive function to print all the nodes at distance k in the tree (or subtree) rooted with given root.
 public static void printKthDown(TreeNode node, int k) {
  if(node == null || k < 0) {  //base case
   return;
  }
  
  if(k == 0) {  //If we reach a k distant node, print it
   System.out.println(node.val);
   return;
  }
  
  // Recur for left and right subtrees
  printKthDown(node.left, k - 1);
  printKthDown(node.right, k - 1);
 }
 
 // Prints all nodes at distance k from a given target node.
 // The k distant nodes may be upward or downward.  This function
 // Returns distance of root from target node, it returns -1 if target
 // node is not present in tree rooted with root.
 public static int printKthDistanceNode(TreeNode root, TreeNode target, int k) {
  if(root == null) {  // Base Case 1: If tree is empty, return -1
   return -1;
  }
  
  // If target is same as root.  Use the downward function
     // to print all nodes at distance k in subtree rooted with
     // target or root
  if(root == target) {
   printKthDown(root, k);
   return 0;
  }
  
  int leftDistance = printKthDistanceNode(root.left, target, k);  // Recur for left subtree
  
  if(leftDistance != -1) {
   // If root is at distance k from target, print root
         // Note that leftDistance is Distance of root's left child from target
   if(leftDistance + 1 == k) {
    System.out.println(root.val);
   }else {  //go to right subtree and print all k-leftDistance-2 distant nodes. Note that the right child is 2 edges away from left child
    printKthDown(root.right, k - leftDistance - 2);
   }
   return leftDistance + 1; // Add 1 to the distance and return value for parent calls
  }
  
  int rightDistance = printKthDistanceNode(root.right, target, k);
  
  if(rightDistance != -1) {
   if(rightDistance + 1 == k) {
    System.out.println(root.val);
   }else {
    printKthDown(root.left, k - rightDistance - 2);
   }
   return rightDistance + 1;
  }
  
  return -1;  // If target was neither present in left nor in right subtree
 }
 
 //create a new binary tree node
 TreeNode newNode(int data) {
  TreeNode temp = new TreeNode(data);
  temp.val = data;
  temp.left = temp.right = null;
  return temp;
 }
 
 public static void main(String[] args) {
  PrintAllNodeAtDistanceKFromANode p = new PrintAllNodeAtDistanceKFromANode();
  TreeNode root = p.newNode(20);
     root.left = p.newNode(8);
     root.right = p.newNode(22);
     root.left.left = p.newNode(4);
     root.left.right = p.newNode(12);
     root.left.right.left = p.newNode(10);
     root.left.right.right = p.newNode(14);
     
     TreeNode target = root.left.right;
     int k = 2;
     printKthDistanceNode(root, target, k);
     
 }
}

No comments:

Post a Comment