L30. Print all the Nodes at a distance of K in Binary Tree

1. Why Do We Need This? (Problem Statement & Status Quo)

Binary trees are inherently unidirectional data structures. Standard tree traversal algorithms allow us to move downward from a root node to its left and right children, but they do not provide built-in upward pointers to navigate back to parent nodes.

When presented with a specific Target Node in a binary tree and asked to find and return all nodes situated at an exact distance $K$ (measured in edges), standard traversals fall short.

The target nodes can be located anywhere: upward toward the ancestors, or downward in different subtrees branching off those ancestors.
Because we cannot natively travel upward from a child node to its parent, we encounter a structural bottleneck: how do we traverse in all directions (left child, right child, and parent) uniformly?

2. Core Mental Model & Real-Life Analogy

Real-Life Analogy: The City Grid and Fire Hydrants

Imagine you are standing at a specific address (the Target Node) in a city, and you need to dispatch emergency personnel to every building exactly $K$ blocks away from you.

You cannot just walk down the streets (downward child pointers); you must also be able to walk up the avenues toward the cross streets and back alleys (parent pointers).
To solve this, you first map out the neighborhood so every building knows its immediate connecting street and parent block.
Then, you dispatch search teams outward in all four directions simultaneously (radially outward: Up, Left, Right), expanding outward block-by-block (distance 1, distance 2, ..., up to distance $K$ ). Once your search radius hits $K$ , every building your teams reach at that exact moment forms your answer set.

CODE

          [Target Building (5)]
                 /     \
   (Upward)     /       \   (Downward Left / Right)
    Parent     Left      Right
  [Node 3]   [Node 6]   [Node 2]

3. The Two-Step Solution Architecture

To conquer the unidirectional limitation of binary trees, the algorithm is split into two distinct phases:

Phase 1: Parent Mapping (The Setup)
- Perform a Breadth-First Search (BFS) or Depth-First Search (DFS) traversal across the entire binary tree.
- Store every node's parent reference inside a Hash Map (parent_map). This equips every node with an upward connection.
Phase 2: Radial BFS Traversal (The Core Search)
- Starting from the target node, initiate a level-order traversal (BFS) using a Queue.
- Maintain a visited Hash Set to ensure nodes are never processed twice (preventing infinite loops between parents and children).
- Expand radially outward: check Left Child, Right Child, and Parent.
- Increment the distance level-by-level until the current distance equals $K$ . The nodes remaining in the queue at distance $K$ are our exact solution.

4. Phase 1: Step-by-Step Parent Pointer Mapping

Before searching for nodes at distance $K$ , we must construct our roadmap. Using a standard level-order traversal (BFS), we populate a hash map that links every child node to its parent.

Structure & Concept Flow

Queue initialization: [3 (root)]

Pop 3:

- Left child is 5 -> Record parent[5] = 3, push 5 to queue

- Right child is 1 -> Record parent[1] = 3, push 1 to queue

Queue state: [5, 1]

Pop 5:

- Left child is 6 -> Record parent[6] = 5, push 6 to queue

- Right child is 2 -> Record parent[2] = 5, push 2 to queue

...

Code Implementation: Parent Mapping

C++

// C++ Implementation for Mapping Parents
void markParents(TreeNode* root, unordered_map<TreeNode*, TreeNode*>& parent_track) {
    if (!root) return;
    queue<TreeNode*> q;
    q.push(root);
    
    while (!q.empty()) {
        TreeNode* current = q.front();
        q.pop();
        
        if (current->left) {
            parent_track[current->left] = current;
            q.push(current->left);
        }
        if (current->right) {
            parent_track[current->right] = current;
            q.push(current->right);
        }
    }
}

JAVA

// Java Implementation for Mapping Parents
private void markParents(TreeNode root, Map<TreeNode, TreeNode> parent_track) {
    if (root == null) return;
    Queue<TreeNode> queue = new LinkedList<>();
    queue.offer(root);
    
    while (!queue.isEmpty()) {
        TreeNode current = queue.poll();
        
        if (current.left != null) {
            parent_track.put(current.left, current);
            queue.offer(current.left);
        }
        if (current.right != null) {
            parent_track.put(current.right, current);
            queue.offer(current.right);
        }
    }
}

5. Phase 2: Radial BFS Traversal & Distance Tracking

Once the parent map is established, we initiate our radial search from the target node. We use a queue for level-order processing, a visited hash map to track processed nodes, and a current_level counter.

Detailed Walkthrough (Target = 5, K = 2)

Initial State:
- Queue: [5]
- Visited: {[5]: true}
- current_level = 0
Iteration 1 (current_level = 0 $\rightarrow$ 1):
- Dequeue 5. Check its three directional neighbors:
  - Parent (3): Not visited $\rightarrow$ mark visited, push to queue.
  - Left Child (6): Not visited $\rightarrow$ mark visited, push to queue.
  - Right Child (2): Not visited $\rightarrow$ mark visited, push to queue.
- Queue state: [3, 6, 2]
- Increment current_level to 1.
Iteration 2 (current_level = 1 $\rightarrow$ 2):
- Process all nodes at level 1 (3, 6, 2) together:
  - From 3: Parent is null. Left is 5 (visited). Right is 1 (not visited $\rightarrow$ mark visited, push).
  - From 6: Parent is 5 (visited). No children.
  - From 2: Parent is 5 (visited). Left is 7 (not visited $\rightarrow$ mark visited, push). Right is 4 (not visited $\rightarrow$ mark visited, push).
- Queue state: [1, 7, 4]
- Increment current_level to 2.
Termination Check:
- Since current_level has reached our target distance $K$ ( $2$ ), we halt further expansion.
- The nodes currently residing in the queue ([1, 7, 4]) are precisely at a distance of 2 from node 5.

6. Complete Production-Ready Code

C++ Implementation

C++

#include <vector>
#include <queue>
#include <unordered_map>

using namespace std;

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution {
private:
    void markParents(TreeNode* root, unordered_map<TreeNode*, TreeNode*>& parent_track) {
        queue<TreeNode*> q;
        q.push(root);
        while(!q.empty()) {
            TreeNode* current = q.front();
            q.pop();
            if(current->left) {
                parent_track[current->left] = current;
                q.push(current->left);
            }
            if(current->right) {
                parent_track[current->right] = current;
                q.push(current->right);
            }
        }
    }

public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        unordered_map<TreeNode*, TreeNode*> parent_track;
        markParents(root, parent_track);
        
        unordered_map<TreeNode*, bool> visited;
        queue<TreeNode*> q;
        
        q.push(target);
        visited[target] = true;
        
        int current_level = 0;
        while(!q.empty()) {
            int size = q.size();
            if(current_level++ == k) break;
            
            for(int i = 0; i < size; i++) {
                TreeNode* current = q.front();
                q.pop();
                
                // Check Left Child
                if(current->left && !visited[current->left]) {
                    q.push(current->left);
                    visited[current->left] = true;
                }
                // Check Right Child
                if(current->right && !visited[current->right]) {
                    q.push(current->right);
                    visited[current->right] = true;
                }
                // Check Parent Upward
                if(parent_track[current] && !visited[parent_track[current]]) {
                    q.push(parent_track[current]);
                    visited[parent_track[current]] = true;
                }
            }
        }
        
        vector<int> result;
        while(!q.empty()) {
            TreeNode* node = q.front();
            q.pop();
            result.push_back(node->val);
        }
        return result;
    }
};

Java Implementation