Open In App

Print nodes at k distance from root | Iterative

Last Updated : 01 Oct, 2024
Comments
Improve
Suggest changes
Like Article
Like
Report

Given a binary tree, and an integer k. The task is to return all the nodes which are at kdistance from the root. 

Example :

Input:

Print-nodes-at-k-distance-from-root

Output: 2 9 13
Explanation: In the above tree 2, 9 & 13 are at distance 2 from root. 

Input:

Print-nodes-at-k-distance-from-root-2

Output: 5 11
Explanation: In the above tree 5 & 11 are at distance 1 from root. 

[Expected Approach - 1] Using Queue - O(n) time and O(n) Space

This idea is to use level order traversal to find all nodes at distancekfrom the root. Start with an empty queue, enqueue the root, and set the level to 0. While the queue isn't empty, if the level equals k, store and return node values. For each node, dequeue it and enqueue its children. Increment the level after processing all nodes. return an empty list if no nodes are at distance k.

Follow the steps below to solve the problem:

  • Start with an empty queue and enqueue the root node. Set the current level (lvl) to 0.
  • While the queue is not empty, repeat the following steps:
    • Determine the number of nodes at the current level (n).
    • If the current level equals k, store the values of the nodes in the result vector by dequeue nodes from the queue.
    • For each node at the current level, dequeue it and enqueue its left and right children if they exist.
    • After processing all nodes at the current level, increment the level counter.
  • If no nodes are found at distance k, return an empty list.

Below is the implementation of the above approach: 

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

class Node {
public:
	int data;
	Node *left, *right;
	Node(int x) {
     	data = x;
      	left = right = nullptr;
    }
};

// Function to find all nodes at distance k from the root
vector<int> Kdistance(Node* root, int k) {
  
    // If the root is NULL, return an empty vector
    if (root == nullptr)
        return {};

    vector<int> result;

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

    while (!q.empty()) {
      
        // Get the number of nodes at the current level
        int n = q.size();

        if (lvl == k) {
          
            // Collect all nodes at this level
            for (int i = 0; i < n; i++) {
                Node* temp = q.front();
                result.push_back(temp->data);
                q.pop();
            }
            // Return the result as we've found 
          	// nodes at distance k
            return result;
        }

        // Process all nodes at the current level
        for (int i = 0; i < n; i++) {
            Node* temp = q.front();
            q.pop();
          
            // If the left child exists, add it to the queue
            if (temp->left != nullptr)
                q.push(temp->left);
          
            // If the right child exists, add it to the queue
            if (temp->right != nullptr)
                q.push(temp->right);
        }

        // Move to the next level
        lvl += 1;
    }

    return result;
}

int main() {
  
	// Constructed binary tree:
    //         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);

	int k = 2;
	vector<int> result = Kdistance(root, k);
	for(int i : result) {
     	cout << i << " "; 
    }

	return 0;
}
Java
// Java code to implement the iterative 
// approach using a Queue

import java.util.*;

class Node {
    int data;
    Node left, right;

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

class GfG {
  
    // Function to find all nodes at distance k from the root
    static List<Integer> Kdistance(Node root, int k) {
      
        // If the root is NULL, return an empty list
        if (root == null)
            return new ArrayList<>();

        List<Integer> result = new ArrayList<>();

        // Create an empty queue for level order traversal
        Queue<Node> q = new LinkedList<>();
        q.add(root);
        int lvl = 0;
        while (!q.isEmpty()) {
          
            // Get the number of nodes at the current level
            int n = q.size();

            if (lvl == k) {
              
                // Collect all nodes at this level
                for (int i = 0; i < n; i++) {
                    result.add(q.peek().data);
                    q.poll();
                }
              
                // Return the result as we've found
              	// nodes at distance k
                return result;
            }

            // Process all nodes at the current level
            for (int i = 0; i < n; i++) {
                Node temp = q.poll();
              
                // If the left child exists, add it to the queue
                if (temp.left != null)
                    q.add(temp.left);
              
                // If the right child exists, add it to the queue
                if (temp.right != null)
                    q.add(temp.right);
            }

            // Move to the next level
            lvl += 1;
        }
        return result;
    }

    public static void main(String[] args) {
      
        // Constructed binary tree:
        //         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);

        List<Integer> result = Kdistance(root, 2);
        for (int i : result) System.out.print(i + " ");
    }
}
Python
# Python code to implement the iterative 
# approach using a Queue

from collections import deque

class Node:
    def __init__(self, x):
        self.data = x
        self.left = None
        self.right = None

def Kdistance(root, k):
  
    # If the root is NULL, return an empty list
    if root is None:
        return []
    
    result = []
    
    # Create an empty queue for level order traversal
    q = deque([root])
    lvl = 0
    while q:
      
        # Get the number of nodes at the current level
        n = len(q)
        
        if lvl == k:
          
            # Collect all nodes at this level
            for _ in range(n):
                result.append(q.popleft().data)
                
            # Return the result as we've found nodes at distance k
            return result
        
        # Process all nodes at the current level
        for _ in range(n):
            temp = q.popleft()
            
            # If the left child exists, add it to the queue
            if temp.left:
                q.append(temp.left)
                
            # If the right child exists, add it to the queue
            if temp.right:
                q.append(temp.right)

        # Move to the next level
        lvl += 1
    return result

if __name__ == "__main__":
  
    # Constructed binary tree:
    #         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)

    result = Kdistance(root, 2)
    print(" ".join(map(str, result)))
C#
// C# code to implement the iterative 
// approach using a Queue
using System;
using System.Collections.Generic;

class Node {
    public int data;
    public Node left, right;

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

class GfG {
  
    // Function to find all nodes at distance k from the root
    static List<int> Kdistance(Node root, int k) {
      
        // If the root is NULL, return an empty list
        if (root == null)
            return new List<int>();

        List<int> result = new List<int>();

        // Create an empty queue for level order traversal
        Queue<Node> q = new Queue<Node>();
        q.Enqueue(root);
        int lvl = 0;
        while (q.Count > 0) {
          
            // Get the number of nodes at the 
          	// current level
            int n = q.Count;

            if (lvl == k) {
              
                // Collect all nodes at this level
                for (int i = 0; i < n; i++) {
                    result.Add(q.Peek().data);
                    q.Dequeue();
                }
                // Return the result as we've found 
              	// nodes at distance k
                return result;
            }

            // Process all nodes at the current level
            for (int i = 0; i < n; i++) {
                Node temp = q.Dequeue();
              
                // If the left child exists, add it to the queue
                if (temp.left != null)
                    q.Enqueue(temp.left);
              
                // If the right child exists, add it to the queue
                if (temp.right != null)
                    q.Enqueue(temp.right);
            }

            // Move to the next level
            lvl += 1;
        }
        return result;
    }

    static void Main(string[] args) {
      
        // Constructed binary tree:
        //         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);

        List<int> result = Kdistance(root, 2);
        foreach (int i in result) {
            Console.Write(i + " ");
        }
    }
}
JavaScript
// Javascript code to implement the iterative 
// approach using a Queue

class Node {
    constructor(x) {
        this.data = x;
        this.left = null;
        this.right = null;
    }
}

// Function to find all nodes at distance k 
// from the root
function Kdistance(root, k) {

    // If the root is NULL, return an empty array
    if (root === null)
        return [];

    let result = [];
    
    // Create an empty queue for level order traversal
    let q = [];
    q.push(root);
    let lvl = 0;
    while (q.length > 0) {
    
        // Get the number of nodes at the current level
        let n = q.length;

        if (lvl === k) {
        
            // Collect all nodes at this level
            for (let i = 0; i < n; i++) {
                result.push(q[0].data);
                q.shift();
            }
            
            // Return the result as we've found nodes
            // at distance k
            return result;
        }

        // Process all nodes at the current level
        for (let i = 0; i < n; i++) {
            let temp = q.shift();
            
            // If the left child exists, add it to the queue
            if (temp.left !== null)
                q.push(temp.left);
                
            // If the right child exists, add it to the queue
            if (temp.right !== null)
                q.push(temp.right);
        }

        // Move to the next level
        lvl += 1;
    }
    return result;
}

// Constructed binary tree:
//         1 
//       /   \ 
//      2     3 
//     / \   / 
//    4   5 8 

let 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);

let result = Kdistance(root, 2);
console.log(result.join(" "));

Output
4 5 8 

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

[Expected Approach - 2] Using Stack - O(n) time and O(n) Space

The idea is to use stack data structure to find all nodes at distance k from the root. Initialize a stack with the root node and its level (0). While the stack isn't empty, pop the top pair and process the node. If the node is NULL, skip to the next iteration. If the node's level equals k, add its data to the result. Otherwise, push the right child with an incremented level, followed by the left child. Return the result vector containing all nodes at distance k from the root.

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;

class Node {
public:
	int data;
	Node *left, *right;
	Node(int x) {
     	data = x;
      	left = right = nullptr;
    }
};

// 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> > stk;
	stk.push({root, 0});

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

		if (curr == nullptr) {
			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
		stk.push({curr->right, level + 1});

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

	return result;
}

int main() {
  
	// Constructed binary tree:
    //         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);

	int k = 2;
	vector<int> result = Kdistance(root, k);
	for(int i : result) {
     	cout << i << " "; 
    }

	return 0;
}
Java
// Java code to implement the iterative 
// approach using a stack
import java.util.*;

class Node {
    int data;
    Node left, right;

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

class Pair {
    Node node;
    int level;

    public Pair(Node node, int level) {
        this.node = node;
        this.level = level;
    }
}

class GfG {

    // Function to perform iterative DFS traversal and find
  	// all nodes at distance K
    static List<Integer> Kdistance(Node root, int k) {
        List<Integer> result = new ArrayList<>();
        Stack<Pair> stk = new Stack<>();
        stk.push(new Pair(root, 0));

        while (!stk.isEmpty()) {
            Pair top = stk.pop();
            Node curr = top.node;
            int level = top.level;

            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.add(curr.data);
            }

            // Push the right child onto the stack with its 
          	// level incremented by 1
            stk.push(new Pair(curr.right, level + 1));

            // Push the left child onto the stack with its
          	// level incremented by 1
            stk.push(new Pair(curr.left, level + 1));
        }

        return result;
    }

    public static void main(String[] args) {

        // Constructed binary tree:
        //         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);

        int k = 2;
        List<Integer> result = Kdistance(root, k);

        for (int i : result) {
            System.out.print(i + " ");
        }
    }
}
Python
# Python code to implement the iterative 
# approach using a stack

class Node:
    def __init__(self, x):
        self.data = x
        self.left = None
        self.right = None

# Function to perform iterative DFS traversal 
# and find all nodes at distance K
def Kdistance(root, k):
    result = []
    stk = [(root, 0)]

    while stk:
        curr, level = stk.pop()

        if curr is None:
            continue

        # If the current node is at distance K 
        # from the root,
        # add its data to the result
        if level == k:
            result.append(curr.data)

        # Push the right child onto the stack with 
        # its level incremented by 1
        stk.append((curr.right, level + 1))

        # Push the left child onto the stack with its 
        # level incremented by 1
        stk.append((curr.left, level + 1))

    return result

# Constructed binary tree:
#         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)

k = 2
result = Kdistance(root, k)
print(" ".join(map(str, result)))
C#
// C# code to implement the iterative approach 
// using a stack
using System;
using System.Collections.Generic;

class Node {
    public int data;
    public Node left, right;

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

class GfG {

    // Function to perform iterative DFS traversal and find all
    // nodes at distance K
    static List<int> Kdistance(Node root, int k) {
        List<int> result = new List<int>();
        Stack<Tuple<Node, int>> stk = new Stack<Tuple<Node, int>>();
        stk.Push(new Tuple<Node, int>(root, 0));

        while (stk.Count > 0) {
            Node curr = stk.Peek().Item1;
            int level = stk.Peek().Item2;
            stk.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.Add(curr.data);
            }

            // Push the right child onto the stack with its
            // level incremented by 1
            stk.Push(new Tuple<Node, int>(curr.right, level + 1));

            // Push the left child onto the stack with its level
            // incremented by 1
            stk.Push(new Tuple<Node, int>(curr.left, level + 1));
        }

        return result;
    }

    static void Main() {
      
        // Constructed binary tree:
        //         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);

        int k = 2;
        List<int> result = Kdistance(root, k);
      
        foreach (int i in result) {
            Console.Write(i + " ");
        }
    }
}
JavaScript
// JavaScript code to implement the 
// iterative approach using a stack

class Node {
    constructor(x) {
        this.data = x;
        this.left = this.right = null;
    }
}

// Function to perform iterative DFS traversal 
// and find all nodes at distance K
function Kdistance(root, k) {
    let result = [];
    let stk = [[root, 0]];

    while (stk.length > 0) {
        let [curr, level] = stk.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(curr.data);
        }

        // Push the right child onto the stack with
        // its level incremented by 1
        stk.push([curr.right, level + 1]);

        // Push the left child onto the stack with its
        // level incremented by 1
        stk.push([curr.left, level + 1]);
    }

    return result;
}

// Constructed binary tree:
//         1
//       /   \
//      2     3
//     / \   /
//    4   5 8

let 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);

let k = 2;
let result = Kdistance(root, k);
console.log(result.join(" "));

Output
4 5 8 

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

Releted articles:


Print nodes at k distance from root iteratively
Article Tags :
Practice Tags :

Similar Reads