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.

Iterative approach using a stack:

Follow the below steps to implement the approach:

• Initializing a stack with a pair of the root node and its level, which is 0.
• While the stack is not empty, pop a pair from the top of the stack and process its node.
  • If the current node is NULL, simply continue to the next iteration.
  • If the current node’s level is equal to K, add its data to the result.
  • Otherwise, push the current node’s right child onto the stack with its level incremented by 1, followed by the current node’s left child with its level incremented by 1.
• Finally, return the result vector containing all the nodes at a distance K from the root.

Below is the implementation of the above approach:

C++

// C++ code to implement the iterative approach using a
// stack
#include <bits/stdc++.h>
using namespace std;

// Node structure
struct Node {
    int data;
    Node *left, *right;
};

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

// Function to perform iterative DFS traversal and find all
// nodes at distance K
vector<int> Kdistance(struct Node* root, int k)
{
    vector<int> result;
    stack<pair<Node*, int> > s;
    s.push(make_pair(root, 0));

    while (!s.empty()) {
        Node* curr = s.top().first;
        int level = s.top().second;
        s.pop();

        if (curr == NULL) {
            continue;
        }

        // If the current node is at distance K from the
        // root, add its data to the result
        if (level == k) {
            result.push_back(curr->data);
        }

        // Push the right child onto the stack with its
        // level incremented by 1
        s.push(make_pair(curr->right, level + 1));

        // Push the left child onto the stack with its level
        // incremented by 1
        s.push(make_pair(curr->left, level + 1));
    }

    return result;
}

// Driver code
int main()
{
    // Create the binary tree
    Node* root = newNode(10);
    root->left = newNode(10);
    root->right = newNode(4);
    root->left->left = newNode(4);
    root->left->right = newNode(10);

    int k = 1;
    vector<int> result = Kdistance(root, k);

    // Print the nodes at distance K from the root
    if (result.empty()) {
        cout << "No nodes found at distance " << k
             << " from the root.\n";
    }
    else {
        cout << "Nodes at distance " << k
             << " from the root: ";
        for (int i = 0; i < result.size(); i++) {
            cout << result[i] << " ";
        }
        cout << endl;
    }

    return 0;
}
// This code is contributed by Veerendra_Singh_Rajpoot

Output

Nodes at distance 1 from the root: 10 4 

Time Complexity: O(N) , The time complexity of the iterative approach using a stack to find all nodes at a distance K from the root of a binary tree is O(N), where N is the number of nodes in the tree. This is because we need to visit every node in the tree exactly once to determine its level and add it to the result vector if its level is equal to K.

Auxiliary Space: O(N) , The auxiliary space complexity of this approach is also O(N), where N is the number of nodes in the tree. This is because we need to store information about each node in the stack until we process it, which requires O(N) space in the worst case for a completely unbalanced tree. 

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