Print nodes at k distance from root

Given a root of a tree, and an integer k. Print all the nodes which are at k distance from root.

For example, in the below tree, 4, 5 & 8 are at distance 2 from root.

Recommended: Please solve it on “PRACTICE” first, before moving on to the solution.

The problem can be solved using recursion. Thanks to eldho for suggesting the solution.

C

#include <stdio.h>
#include <stdlib.h>

/* A binary tree node has data, pointer to left child
   and a pointer to right child */
struct node
{
   int data;
   struct node* left;
   struct node* right;
};

void printKDistant(struct node *root , int k)    
{
   if(root == NULL) 
      return;
   if( k == 0 )
   {
      printf( "%d ", root->data );
      return ;
   }
   else
   {      
      printKDistant( root->left, k-1 ) ;
      printKDistant( root->right, k-1 ) ;
   }
}

/* Helper function that allocates a new node with the
   given data and NULL left and right pointers. */
struct node* newNode(int data)
{
  struct node* node = (struct node*)
                       malloc(sizeof(struct node));
  node->data = data;
  node->left = NULL;
  node->right = NULL;

  return(node);
}

/* Driver program to test above functions*/
int main()
{

  /* Constructed binary tree is
            1
          /   \
        2      3
      /  \    /
    4     5  8 
  */
  struct node *root = newNode(1);
  root->left        = newNode(2);
  root->right       = newNode(3);
  root->left->left  = newNode(4);
  root->left->right = newNode(5);
  root->right->left = newNode(8);  

  printKDistant(root, 2);

  getchar();
  return 0;
}

Java

// Java program to print nodes at k distance from root
 
/* A binary tree node has data, pointer to left child
   and a pointer to right child */
class Node 
{
    int data;
    Node left, right;
 
    Node(int item) 
    {
        data = item;
        left = right = null;
    }
}
 
class BinaryTree 
{
    Node root;
 
    void printKDistant(Node node, int k) 
    {
        if (node == null)
            return;
        if (k == 0) 
        {
            System.out.print(node.data + " ");
            return;
        } 
        else 
        {
            printKDistant(node.left, k - 1);
            printKDistant(node.right, k - 1);
        }
    }
    
    /* Driver program to test above functions */
    public static void main(String args[]) {
        BinaryTree tree = new BinaryTree();
        
        /* Constructed binary tree is
                1
              /   \
             2     3
            /  \   /
           4    5 8 
        */
        tree.root = new Node(1);
        tree.root.left = new Node(2);
        tree.root.right = new Node(3);
        tree.root.left.left = new Node(4);
        tree.root.left.right = new Node(5);
        tree.root.right.left = new Node(8);
 
        tree.printKDistant(tree.root, 2);
    }
}
 
// This code has been contributed by Mayank Jaiswal

Python


# Python program to find the nodes at k distance from root

# A Binary tree node
class Node:
    
    # Constructor to create a new node
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None

def printKDistant(root, k):
    
    if root is None:
        return 
    if k == 0:
        print root.data,
    else:
        printKDistant(root.left, k-1)
        printKDistant(root.right, k-1)

# Driver program to test above function
"""
   Constructed binary tree is
            1
          /   \
        2      3
      /  \    /
    4     5  8 
"""
root = Node(1)
root.left = Node(2)
root.right = Node(3)
root.left.left = Node(4)
root.left.right = Node(5)
root.right.left = Node(8)

printKDistant(root, 2)

# This code is contributed by Nikhil Kumar Singh(nickzuck_007)

The above program prints 4, 5 and 8.

Time Complexity: O(n) where n is number of nodes in the given binary tree.