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. 

            1
          /   \
        2      3
      /  \    /
    4     5  8 

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

Implementation:

C++

#include<bits/stdc++.h>

using namespace std;

/* A binary tree node has data,
pointer to left child and
a pointer to right child */
class node
{
    public:
    int data;
    node* left;
    node* right;

    /* Constructor that allocates a new node with the
    given data and NULL left and right pointers. */
    node(int data)
    {
        this->data = data;
        this->left = NULL;
        this->right = NULL;
    }
};

void printKDistant(node *root , int k)
{
    if(root == NULL|| k < 0 )
        return;
    if( k == 0 )
    {
        cout << root->data << " ";
         return;
    }

        printKDistant( root->left, k - 1 ) ;
        printKDistant( root->right, k - 1 ) ;

}


/* Driver code*/
int main()
{

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

    printKDistant(root, 2);
    return 0;
}

// This code is contributed by rathbhupendra

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|| k < 0 )
      return;
   if( k == 0 )
   {
      printf( "%d ", root->data );
      return ;
   }

      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|| k < 0 )
              //Base case
            return;
        if (k == 0)
        {
            System.out.print(node.data + " ");
            return;
        }
       //recursively traversing
            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

Python3

# 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,end=' ')
    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)

C#

using System;

// c# 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 */
public class Node
{
    public int data;
    public Node left, right;

    public Node(int item)
    {
        data = item;
        left = right = null;
    }
}

public class BinaryTree
{
    public Node root;

    public virtual void printKDistant(Node node, int k)
    {
        if (node == null|| k < 0 )
        {
            return;
        }
        if (k == 0)
        {
            Console.Write(node.data + " ");
            return;
        }

            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 is contributed by Shrikant13

Javascript

<script>

// Javascript 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
{
    constructor(item)
    {
        this.data = item;
        this.left = null;
        this.right = null;
    }
}

var root =null;

function printKDistant(node, k)
{
    if (node == null|| k < 0 )
    {
        return;
    }
    if (k == 0)
    {
        document.write(node.data + " ");
        return;
    }

        printKDistant(node.left, k - 1);
        printKDistant(node.right, k - 1);

}


/* Driver program to test above functions */

/* Constructed binary tree is
        1
      /   \
     2     3
    /  \   /
   4    5 8
*/
root = new Node(1);
root.left = new Node(2);
root.right = new Node(3);
root.left.left = new Node(4);
root.left.right = new Node(5);
root.right.left = new Node(8);
printKDistant(root, 2);

// This code is contributed by importantly.
</script>

Output

4 5 8 

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

Space Complexity : O(height of the binary tree).

Note-

• If it’s true print the node – Always check the K distance == 0 at every node
• the left or right subtree – Decrement the distance by 1 when you are passing to its subtree 

Another Approach– We can do a level order traversal and keep track of the level.when current level is equal to k, then print all the nodes of that level.

Python3

# 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):
   # check if root is None
    if root is None:
        return
    q = []
    # ans = []
    q.append(root)
    lvl = 0  # tracking of level
    while(q):
        n = len(q)
        # when lvl becomes k we add all values of q in ans.
        if lvl == k:

            for i in range(n):
                print((q[i].data), end=" ")
            return

        for i in range(1, n+1):
            temp = q.pop(0)
            if temp.left:
                q.append(temp.left)
            if temp.right:
                q.append(temp.right)
        lvl += 1
    # if after traversing ,if lvl is less than k ,
    # that means nodes at k distance does not exist.
    if lvl < k:
        return

# 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 Vivek Maddeshiya

Output

4 5 8 

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

Please write comments if you find the above code/algorithm incorrect, or find better ways to solve the same problem.