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.

C++

// C++ Program to find the nodes at k distance from root

#include <bits/stdc++.h>
using namespace std;

// A Binary Tree Node
struct Node {
  int data;
  struct Node *left, *right;
};

void printKDistant(Node* root, int k)
{
  // Base Case
  if (root == NULL)
    return;

  // Create an empty queue for level order traversal
  queue<Node*> q;
  q.push(root);
  int lvl = 0;

  while (!q.empty()) {
    int n = q.size();

    if (lvl == k) {
      for (int i = 0; i < n; i++) {
        cout << q.front()->data << " ";
        q.pop();
      }
      return;
    }

    for (int i = 0; i < n; i++) {
      Node* temp = q.front();
      q.pop();
      if (temp->left != NULL)
        q.push(temp->left);
      if (temp->right != NULL)
        q.push(temp->right);
    }
    lvl += 1;
  }
}

// Utility function to create a new tree node
Node* newNode(int data)
{
  Node* temp = new Node;
  temp->data = data;
  temp->left = temp->right = NULL;
  return temp;
}

// Driver program to test above functions
int main()
{
  /* Constructed binary tree is
             1
           /   \
         2      3
       /  \    /
     4     5  8
    */

  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);
  return 0;
}

// This is code is contributed by Yash
// Agarwal(yashagarwal2852002)

Java

// Java Program to find the nodes at k distance from root
import java.util.LinkedList;
import java.util.Queue;

/* Class to represent Tree node */
class Node {
  int data;
  Node left, right;

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

/* Class to print K distance Element */
class BinaryTree {

  Node root;

  void printKDistant(Node node, int k)
  {
    Queue<Node> queue = new LinkedList<Node>();
    queue.add(root);
    int lvl = 0;
    while (!queue.isEmpty()) {
      int n = queue.size();

      if (lvl == k) {
        for (int i = 0; i < n; i++) {
          Node tempNode = queue.poll();
          System.out.print(tempNode.data + " ");
        }
        return;
      }

      for (int i = 0; i < n; i++) {
        Node tempNode = queue.poll();
        if (tempNode.left != null)
          queue.add(tempNode.left);
        if (tempNode.right != null)
          queue.add(tempNode.right);
      }
      lvl++;
    }
  }

  public static void main(String args[])
  {
    /* Constructed binary tree is
            1
          /   \
        2      3
      /  \    /
    4     5  8
*/
    BinaryTree tree = new BinaryTree();
    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 is code is contributed by Yash
// Agarwal(yashagarwal2852002)

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

C#

// C# program to find the nodes at k distance from root
using System;
using System.Collections.Generic;

// Class to represent Tree node
public class Node {
  public int data;
  public Node left, right;

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

// Class to print K distance elements from the root
public class GFG {

  public Node root;

  public void printKDistant(Node node, int k)
  {
    Queue<Node> queue = new Queue<Node>();
    queue.Enqueue(node);
    int lvl = 0;

    while (queue.Count != 0) {
      int n = queue.Count;

      if (lvl == k) {
        for (int i = 0; i < n; i++) {
          Node tempNode = queue.Dequeue();
          Console.Write(tempNode.data + " ");
        }
        return;
      }

      for (int i = 0; i < n; i++) {
        Node tempNode = queue.Dequeue();
        if (tempNode.left != null)
          queue.Enqueue(tempNode.left);
        if (tempNode.right != null)
          queue.Enqueue(tempNode.right);
      }
      lvl++;
    }
  }

  // Driver code
  public static void Main()
  {
    /* Constructed binary tree is
                  1
                /   \
              2      3
            /  \    /
          4     5  8
        */
    GFG tree = new GFG();
    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 is code is contributed by Yash
// Agarwal(yashagarwal2852002)

Javascript

// Javascript Program to find
// the nodes at K distance from the root
class Node {
    constructor(val) {
        this.data = val;
        this.left = null;
        this.right = null;
    }
}

function printKDistant(root, k)
{

    // Base Case
    if(root == null) return;

    // Create an empty queue for level order traversal
    var queue = [];
    queue.push(root);
    var lvl = 0;

    while (queue.length != 0) {
        var n = queue.length;

        if(lvl == k){
            for(var i = 0; i<n; i++){
                var tempNode = queue.shift();
                console.log(tempNode.data + " ");
            }
            return;
        }

        for(var i=0; i<n; i++){
            var tempNode = queue.shift();
            if(tempNode.left != null) queue.push(tempNode.left);
            if(tempNode.right != null) queue.push(tempNode.right);
        }
        lvl++;
    }
}
    /* 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 Yash Agarwal(yashagarwal2852002)

Output

4 5 8 

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

Space 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.