Print all nodes at distance K from given node: Iterative Approach

Given a binary tree, a target node and an integer K, the task is to find all the nodes that are at distance K from the given target node. 

 

Consider the above-given Tree, For the target node 12. 
Input: K = 1 
Output: 8 10 14

Input: K = 2 
Output: 4 20



Input: K = 3 
output: 22 

Approach: 
There are generally two cases for the nodes at a distance of K: 

  1. Node at a distance K is a child node of the target node.
  2. Node at a distance K is the ancestor of the target node.

The idea is to store the parent node of every node in a hash-map with the help of the Level-order traversal on the tree. Then, Simply Traverse the nodes from the Target node using Breadth-First Search on the left-child, right-child, and the parent node. At any instant when the distance of a node the from the target node is equal to K then print all the nodes of the queue.

Below is the implementation of the above approach: 

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation to print all
// the nodes from the given target
// node iterative approach
 
#include <bits/stdc++.h>
 
using namespace std;
 
// Structure of the Node
struct Node {
    int val;
    Node *left, *right;
};
 
// Map to store the parent
// node of every node of
// the given Binary Tree
unordered_map<Node*, Node*> um;
 
// Functiom to store all nodes
// parent in unordered_map
void storeParent(Node* root)
{
 
    // Make a queue to do level-order
    // Traversal and store parent
    // of each node in unordered map
    queue<Node*> q;
    q.push(root);
     
    // Loop to iterate until the
    // queue is not empty
    while (!q.empty()) {
        Node* p = q.front();
        q.pop();
         
        // Condition if the node is a
        /// root node that storing its
        // parent as NULL
        if (p == root) {
            um[p] = NULL;
        }
         
        // if left child exist of
        // popped out node then store
        // parent as value and node as key
        if (p->left) {
            um[p->left] = p;
            q.push(p->left);
        }
        if (p->right) {
            um[p->right] = p;
            q.push(p->right);
        }
    }
}
 
// Function to find the nodes
// at distance K from give node
void nodeAtDistK(Node* root,
           Node* target, int k)
{
    // Keep track of each node
    // which are visited so that
    // while doing BFS we will
    // not traverse it again
    unordered_set<Node*> s;
    int dist = 0;
    queue<Node*> q;
    q.push(target);
    s.insert(target);
     
    // Loop to iterate over the nodes
    // until the queue is not empty
    while (!q.empty()) {
 
        // if distance is equal to K
        // then we found a node in tree
        // which is distance K
        if (dist == k) {
            while (!q.empty()) {
                cout << q.front()->val << " ";
                q.pop();
            }
        }
 
        // BFS on node's left,
        // right and parent node
        int size = q.size();
        for (int i = 0; i < size; i++) {
            Node* p = q.front();
            q.pop();
 
            // if the left of node is not
            // visited yet then push it in
            // queue and insert it in set as well
            if (p->left &&
                s.find(p->left) == s.end()) {
                q.push(p->left);
                s.insert(p->left);
            }
 
            // if the right of node is not visited
            // yet then push it in queue
            // and insert it in set as well
            if (p->right &&
                s.find(p->right) == s.end()) {
                q.push(p->right);
                s.insert(p->right);
            }
 
            // if the parent of node is not visited
            // yet then push it in queue and
            // insert it in set as well
            if (um[p] && s.find(um[p]) == s.end()) {
                q.push(um[p]);
                s.insert(um[p]);
            }
        }
        dist++;
    }
}
 
// Function to create a newnode
Node* newnode(int val)
{
    Node* temp = new Node;
    temp->val = val;
    temp->left = temp->right = NULL;
    return temp;
}
 
// Driver Code
int main()
{
    Node* root = newnode(20);
    root->left = newnode(8);
    root->right = newnode(22);
    root->right->left = newnode(5);
    root->right->right = newnode(8);
    root->left->left = newnode(4);
    root->left->left->left = newnode(25);
    root->left->right = newnode(12);
    root->left->right->left =
                   newnode(10);
    root->left->right->left->left =
                   newnode(15);
    root->left->right->left->right =
                   newnode(18);
    root->left->right->left->right->right =
                   newnode(23);
    root->left->right->right =
                   newnode(14);
    Node* target = root->left->right;
    storeParent(root);
    nodeAtDistK(root, target, 3);
    return 0;
}
chevron_right

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation to print all
// the nodes from the given target
// node iterative approach
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.Queue;
 
class GFG{
 
// Structure of the Node
static class Node
{
    int val;
    Node left, right;
 
    public Node(int val)
    {
        this.val = val;
        this.left = this.right = null;
    }
};
 
// Map to store the parent
// node of every node of
// the given Binary Tree
static HashMap<Node, Node> um = new HashMap<>();
 
// Functiom to store all nodes
// parent in unordered_map
static void storeParent(Node root)
{
     
    // Make a queue to do level-order
    // Traversal and store parent
    // of each node in unordered map
    Queue<Node> q = new LinkedList<>();
    q.add(root);
 
    // Loop to iterate until the
    // queue is not empty
    while (!q.isEmpty())
    {
        Node p = q.poll();
 
        // Condition if the node is a
        /// root node that storing its
        // parent as NULL
        if (p == root)
        {
            um.put(p, null);
        }
 
        // if left child exist of
        // popped out node then store
        // parent as value and node as key
        if (p.left != null)
        {
            um.put(p.left, p);
            q.add(p.left);
        }
        if (p.right != null)
        {
            um.put(p.right, p);
            q.add(p.right);
        }
    }
}
 
// Function to find the nodes
// at distance K from give node
static void nodeAtDistK(Node root,
                        Node target, int k)
{
     
    // Keep track of each node
    // which are visited so that
    // while doing BFS we will
    // not traverse it again
    HashSet<Node> s = new HashSet<>();
    int dist = 0;
     
    Queue<Node> q = new LinkedList<>();
    q.add(target);
    s.add(target);
 
    // Loop to iterate over the nodes
    // until the queue is not empty
    while (!q.isEmpty())
    {
         
        // If distance is equal to K
        // then we found a node in tree
        // which is distance K
        if (dist == k)
        {
            while (!q.isEmpty())
            {
                System.out.print(q.peek().val + " ");
                q.poll();
            }
        }
 
        // BFS on node's left,
        // right and parent node
        int size = q.size();
        for(int i = 0; i < size; i++)
        {
            Node p = q.poll();
 
            // If the left of node is not
            // visited yet then add it in
            // queue and insert it in set as well
            if (p.left != null && !s.contains(p.left))
            {
                q.add(p.left);
                s.add(p.left);
            }
 
            // If the right of node is not visited
            // yet then add it in queue
            // and insert it in set as well
            if (p.right != null && !s.contains(p.right))
            {
                q.add(p.right);
                s.add(p.right);
            }
 
            // If the parent of node is not visited
            // yet then add it in queue and
            // insert it in set as well
            if (um.get(p) != null &&
                !s.contains(um.get(p)))
            {
                q.add(um.get(p));
                s.add(um.get(p));
            }
        }
        dist++;
    }
}
 
// Driver Code
public static void main(String[] args)
{
    Node root = new Node(20);
    root.left = new Node(8);
    root.right = new Node(22);
    root.right.left = new Node(5);
    root.right.right = new Node(8);
    root.left.left = new Node(4);
    root.left.left.left = new Node(25);
    root.left.right = new Node(12);
    root.left.right.left = new Node(10);
    root.left.right.left.left = new Node(15);
    root.left.right.left.right = new Node(18);
    root.left.right.left.right.right = new Node(23);
    root.left.right.right = new Node(14);
     
    Node target = root.left.right;
     
    storeParent(root);
     
    nodeAtDistK(root, target, 3);
}
}
 
// This code is contributed by sanjeev2552
chevron_right

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.





Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.



Improved By : sanjeev2552

Article Tags :
Practice Tags :