Binary Tree: Given a binary tree and a value k, print all the nodes at distance k

Problem Statement:

You are given a binary tree and a target node and a value k.

You need to find all the nodes at a distance k from the given target node.

Example:

Input:

    /*
     *           10
     *         /    \
     *        8      12
     *       / \    /  \
     *      2   9  11   14
     */

target = 8
k = 2

Solution Explanation:

We will use 2 BFS to solve the problem.

The first BSF will build the parent child relationship.

The second will find all the nodes at a distance k from the target node.

Time Complexity: O(1)
Space Complexity: O(1)

Code Solution

#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_map>
#include <queue>

using namespace std;

struct Node 
{
    int data;
    struct Node *left;
    struct Node *right;
};


struct Node* newNode(int data)
{
    struct Node* newNode = new Node;
    newNode->data = data;
    newNode->left = newNode->right = NULL;

    return (newNode);
}

void display_inorder(Node* root)
{
    if (root == NULL)
        return;

    display_inorder(root->left);
    cout << root->data << " ";
    display_inorder(root->right);
}


void insert_inorder(Node* root, vector<int>& nodes) 
{
    if (root == nullptr) 
    {
        return;
    }
  
    insert_inorder(root->left, nodes);  
    nodes.push_back(root->data);          
    insert_inorder(root->right, nodes); 
}

vector<int> solution(Node* root, Node* target, int K) 
{
    unordered_map<Node*, Node*> parent_track; // maps each node to its parent
    unordered_map<Node*, bool> visited; // keep track of nodes that are already explorec

    queue<Node*> queue; // bfs queue
    queue.push(root);

    while(!queue.empty()) /*First BFS to get a track of parent nodes*/
    {  
        Node* current = queue.front(); 
            queue.pop();

        if(current->left) 
        {
            parent_track[current->left] = current;
            queue.push(current->left);
        }
        if(current->right) 
        {
            parent_track[current->right] = current;
            queue.push(current->right);
        }
    }

    /*Second BFS from the target node*/
    queue.push(target);
    visited[target] = true; //mark it as visited
    int curr_level = 0; //keep track of the distance

    while(!queue.empty()) /*Second BFS to go upto K level from target node, using hashtable info */
    { 
        int size = queue.size();
        if(curr_level++ == K) break;
        for(int i=0; i<size; i++) 
        {
            //explore the neighbours
            Node* current = queue.front(); 
            queue.pop();
            
            if(current->left && !visited[current->left]) 
            {
                queue.push(current->left);
                visited[current->left] = true;
            }
            if(current->right && !visited[current->right]) 
            {
                queue.push(current->right);
                visited[current->right] = true;
            }
            if(parent_track[current] && !visited[parent_track[current]]) 
            {
                queue.push(parent_track[current]);
                visited[parent_track[current]] = true;
            }
        }
    }
    vector<int> result;
    while(!queue.empty()) 
    {
        Node* current = queue.front(); 
        queue.pop();
        result.push_back(current->data);
    }
    return result;
}

int main(void)
{
    /*
     *           10
     *         /    \
     *        8      12
     *       / \    /  \
     *      2   9  11   14
     */
    struct Node* root = newNode(10);
    root->left = newNode(8);
    root->right = newNode(12);
    root->left->left = newNode(2);
    root->left->right = newNode(9);
    root->right->left = newNode(11);
    root->right->right = newNode(14);
    
    Node* target = root->left;
    vector<int> result = solution(root, target, 1);
    
    int n = result.size();

    for (int i = 0; i < n; i++) 
    {
        cout << result[i] << " ";
    }
    cout << endl;

    return 0;
}

Output

2 9 10