Why is a visited set required in the BFS approaches?

Once upward movement or bidirectional edges are introduced, the traversal can move from a child to its parent and then back to the same child. The visited set prevents such repeated traversal.

What happens when K = 0?

The target node itself is at distance 0, so the answer contains only the target.

Why is the opposite subtree searched in the recursive approach?

After traversal has moved upward from the target to an ancestor, returning into the same subtree would move back toward the target. Any new candidate must therefore be the ancestor itself or lie inside its opposite subtree.

Which approach is generally preferred in interviews?

The parent mapping + BFS approach is generally the easiest to implement and explain because every node can effectively move to its left child, right child, or parent while BFS directly represents distance in edges.

Print All Nodes at Distance K from a Target Node

Given the root of a binary tree, a target node present in the tree, and an integer K, return the values of all nodes whose distance from the target is exactly K edges.

The required nodes may lie:

inside the target's subtree,
above the target through its ancestors, or
inside a different subtree reached through an ancestor.

The nodes may be returned in any order.

Example 1

Input:
root = [3, 5, 1, 6, 2, 0, 8, null, null, 7, 4], target = 5, K = 2

Output:
[7, 4, 1]

Explanation:
Nodes 7 and 4 are two edges below target 5. Node 1 is reached through the path 5 -> 3 -> 1, which also contains 2 edges.

Example 2

Input:
root = [1], target = 1, K = 0

Output:
[1]

Explanation:
The target node itself is at distance 0, so it is the only node included in the answer.

Practice it by yourself before moving onto the solutionTry now

Approach 1

A normal binary tree allows movement only from a parent to its children. However, nodes at distance K from the target may also require movement upward toward an ancestor.

For example, node 1 in Example 1 is reached from target 5 using:

5 -> 3 -> 1

To allow movement in both directions, every parent-child connection is treated as an undirected edge. An adjacency list is used to store all neighboring nodes so that movement from parent to child as well as child to parent becomes possible.

A queue is then used for BFS starting from the target. A variable currLevel is maintained because it represents the number of edges between the target and the nodes currently being processed.

A visited set is also maintained because converting the tree into an undirected graph introduces cycles such as:

parent -> child -> parent

The set prevents the same node from being processed repeatedly.

Algorithm

An adjacency list is constructed by traversing the binary tree, and every parent-child connection is stored in both directions so that movement upward and downward becomes possible.
A BFS queue is initialized with the target, while a visited set is initialized with the target so that nodes are not revisited through bidirectional edges.
A variable currLevel is initialized to 0 because the target itself lies at distance 0.
While the queue is not empty and currLevel is smaller than K, all nodes belonging to the current BFS level are processed.
For every processed node, each unvisited neighbor is marked as visited and inserted into the queue because that neighbor lies one additional edge away from the target.
After one complete BFS level has been processed, currLevel is increased by 1.
When currLevel becomes K, further expansion is stopped because every node currently remaining in the queue is exactly K edges away from the target.
The values of all nodes remaining in the queue are collected and returned.

Dry Run

All Nodes Distance K in Binary Tree Approach 1 Dry Run .png

Solution

1#include <bits/stdc++.h>
2using namespace std;
3
4struct TreeNode {
  int val;
  TreeNode* left;
  TreeNode* right;
8
  TreeNode(int value) {
      val = value;
      left = nullptr;
      right = nullptr;
  }
14};
15
16class Solution {
17private:
  // Converts each parent-child connection
  // into an undirected graph edge.
  void buildGraph(
      TreeNode* node,
      unordered_map<TreeNode*, vector<TreeNode*>>& graph
  ) {
      if (node == nullptr) {
          return;
      }
27
      // Both directions are stored so traversal
      // can move downward as well as upward.
      if (node->left != nullptr) {
          graph[node].push_back(node->left);
          graph[node->left].push_back(node);
33
          buildGraph(node->left, graph);
      }
36
      if (node->right != nullptr) {
          graph[node].push_back(node->right);
          graph[node->right].push_back(node);
40
          buildGraph(node->right, graph);
      }
  }
44
45public:
  // Returns all nodes exactly K edges
  // away from the target node.
  vector<int> distanceK(
      TreeNode* root,
      TreeNode* target,
      int K
  ) {
      vector<int> answer;
54
      if (root == nullptr) {
          return answer;
      }
58
      unordered_map<TreeNode*, vector<TreeNode*>> graph;
      buildGraph(root, graph);
61
      queue<TreeNode*> nodesQueue;
      unordered_set<TreeNode*> visited;
64
      nodesQueue.push(target);
66
      // The target is marked immediately so
      // bidirectional edges cannot revisit it.
      visited.insert(target);
70
      // currLevel represents the number of edges
      // from the target to the current BFS level.
      int currLevel = 0;
74
      while (!nodesQueue.empty()) {
76
          // Nodes already in the queue are exactly
          // K edges away, so no further expansion is needed.
          if (currLevel == K) {
              break;
          }
82
          int levelSize = nodesQueue.size();
84
          // Exactly one BFS level is processed here.
          for (int i = 0; i < levelSize; i++) {
              TreeNode* node = nodesQueue.front();
              nodesQueue.pop();
89
              for (TreeNode* neighbor : graph[node]) {
                  if (visited.find(neighbor) == visited.end()) {
92
                      // Mark before insertion so the same node
                      // cannot enter the queue through another edge.
                      visited.insert(neighbor);
                      nodesQueue.push(neighbor);
                  }
              }
          }
100
          // Completing one BFS level increases
          // the distance from the target by one.
          currLevel++;
      }
105
      // Every remaining node is exactly
      // K edges away from the target.
      while (!nodesQueue.empty()) {
          answer.push_back(
              nodesQueue.front()->val
          );
          nodesQueue.pop();
      }
114
      return answer;
  }
117};
118
119int main() {
  TreeNode* root = new TreeNode(3);

Complexity Analysis

Let N be the number of nodes in the binary tree.

Time Complexity: O(N), where N is the number of nodes in the binary tree. Every node and every tree edge is processed a constant number of times while the undirected graph is constructed and BFS is performed.

Space Complexity: O(N). The adjacency list stores information for up to N nodes, while the visited set and BFS queue can also contain up to N nodes in the worst case.

Approach 2

Constructing a complete undirected adjacency list stores more information than is actually required.

Every tree node already provides access to its left and right children. The only missing direction is movement from a child to its parent.

Therefore, a mapping named parentTrack is maintained:

child -> parent

The name parentTrack indicates that the structure is used to keep track of each node's parent so that traversal can move upward when required.

Once this mapping is available, every node effectively has at most three possible neighbors:

left child,
right child,
parent.

A nodesQueue is used for BFS from the target, while a visited set prevents movement back to nodes that have already been processed.

A variable currLevel represents the current distance from the target because every BFS level corresponds to one additional edge.

Algorithm

A queue-based level-order traversal is performed from the root, and a map named parentTrack is constructed so that every non-root node can move upward to its parent.
The root is treated separately in the parent mapping because it has no parent and therefore does not need an upward connection.
A nodesQueue is initialized with the target because BFS must expand outward from the target one edge at a time.
A visited set is initialized with the target so that movement through parent and child links cannot cause the same node to be processed again.
A variable currLevel is initialized to 0 because the target is at distance 0 from itself.
At each BFS level, the left child, right child, and mapped parent of every node are considered whenever they exist and have not already been visited.
Every newly discovered neighbor is marked as visited before being inserted into nodesQueue, preventing duplicate insertion through another direction.
After one complete level has been processed, currLevel is increased because all newly queued nodes are one edge farther from the target.
When currLevel becomes K, further traversal is stopped, and the values of all nodes remaining in nodesQueue are returned.

Dry Run

All Nodes Distance K in Binary Tree Approach 2 Dry Run .png

Solution

1#include <bits/stdc++.h>
2using namespace std;
3
4struct TreeNode {
  int val;
  TreeNode* left;
  TreeNode* right;
8
  TreeNode(int value) {
      val = value;
      left = nullptr;
      right = nullptr;
  }
14};
15
16class Solution {
17private:
  // Builds child-to-parent relationships
  // using level-order traversal.
  void buildParentTrack(
      TreeNode* root,
      unordered_map<TreeNode*, TreeNode*>& parentTrack
  ) {
      queue<TreeNode*> nodesQueue;
      nodesQueue.push(root);
26
      // The root has no parent, so no upward
      // mapping is required for the root itself.
      parentTrack[root] = nullptr;
30
      while (!nodesQueue.empty()) {
          TreeNode* node =
              nodesQueue.front();
34
          nodesQueue.pop();
36
          if (node->left != nullptr) {
              parentTrack[node->left] = node;
              nodesQueue.push(node->left);
          }
41
          if (node->right != nullptr) {
              parentTrack[node->right] = node;
              nodesQueue.push(node->right);
          }
      }
  }
48
49public:
  // Uses child, parent, and sibling directions
  // to perform BFS outward from target.
  vector<int> distanceK(
      TreeNode* root,
      TreeNode* target,
      int K
  ) {
      vector<int> answer;
58
      if (root == nullptr) {
          return answer;
      }
62
      // parentTrack provides the only direction
      // missing from a normal binary tree.
      unordered_map<TreeNode*, TreeNode*> parentTrack;
66
      buildParentTrack(
          root,
          parentTrack
      );
71
      // nodesQueue stores nodes belonging
      // to the current and upcoming BFS levels.
      queue<TreeNode*> nodesQueue;
      nodesQueue.push(target);
76
      // visited prevents movement from a child
      // to its parent and back to the same child.
      unordered_set<TreeNode*> visited;
      visited.insert(target);
81
      // currLevel stores the current
      // distance from the target.
      int currLevel = 0;
85
      while (!nodesQueue.empty()) {
87
          // All queued nodes are already at
          // distance K, so expansion is stopped.
          if (currLevel == K) {
              break;
          }
93
          int levelSize =
              nodesQueue.size();
96
          for (int i = 0; i < levelSize; i++) {
              TreeNode* node =
                  nodesQueue.front();
100
              nodesQueue.pop();
102
              // The left child is one edge
              // farther from the current node.
              if (
                  node->left != nullptr &&
                  visited.find(node->left) == visited.end()
              ) {
                  visited.insert(node->left);
                  nodesQueue.push(node->left);
              }
112
              // The right child is processed
              // in the same downward direction.
              if (
                  node->right != nullptr &&
                  visited.find(node->right) == visited.end()
              ) {
                  visited.insert(node->right);
                  nodesQueue.push(node->right);

Complexity Analysis

Let N be the number of nodes in the binary tree.

Time Complexity: O(N), where N is the number of nodes in the binary tree. Building parentTrack requires visiting every node once, and the BFS from the target visits each reachable node at most once.

Space Complexity: O(N). The parentTrack map, visited set, and nodesQueue can each contain information for up to N nodes in the worst case.

Approach 3

The parent map can be avoided completely by using recursion to determine how far the target lies below each ancestor.

Two types of nodes can exist at distance K from the target:

Nodes below the target.
Nodes reached by moving upward to an ancestor and then either selecting that ancestor or moving into its opposite subtree.

A helper is used to collect nodes lying a specific number of edges below a given node.

Another recursive function returns the distance between the current node and the target. A returned value of -1 indicates that the target does not exist in that subtree. Otherwise, the returned value tells an ancestor how many edges below it the target was found.

Suppose an ancestor is distance edges away from the target:

if distance == K, that ancestor itself is included;
otherwise, the ancestor's opposite subtree is searched for nodes that can complete the remaining distance.

This allows both downward and upward paths to be handled without storing parent pointers.

Algorithm

A helper is used to collect nodes that lie exactly a required number of edges below a given node.
A DFS is started from the root, and each recursive call is made to return the distance from the current node to the target, while -1 is returned when the target is absent from that subtree.
When the target node is reached, all descendants exactly K edges below it are collected, and distance 0 is returned to its parent.
If a valid distance is returned from the left child, that value is increased by 1 to obtain the current node's distance from the target.
If this updated distance equals K, the current ancestor is inserted into the answer. Otherwise, the right subtree is searched for nodes whose additional distance completes exactly K edges.
The same symmetric process is performed when the target is found in the right subtree, with the left subtree being treated as the opposite subtree.
The calculated target distance is propagated upward until every relevant ancestor and opposite subtree has been processed.

Dry Run

All Nodes Distance K in Binary Tree Approach 3 Dry Run .png

Solution

1#include <bits/stdc++.h>
2using namespace std;
3
4struct TreeNode {
  int val;
  TreeNode* left;
  TreeNode* right;
8
  TreeNode(int value) {
      val = value;
      left = nullptr;
      right = nullptr;
  }
14};
15
16class Solution {
17private:
  // Collects nodes exactly distance edges
  // below the given subtree root.
  void collectDown(
      TreeNode* node,
      int distance,
      vector<int>& answer
  ) {
      if (
          node == nullptr ||
          distance < 0
      ) {
          return;
      }
31
      if (distance == 0) {
          answer.push_back(node->val);
          return;
      }
36
      collectDown(
          node->left,
          distance - 1,
          answer
      );
42
      collectDown(
          node->right,
          distance - 1,
          answer
      );
  }
49
  // Returns the distance from node to target,
  // or -1 when target is absent from the subtree.
  int findTarget(
      TreeNode* node,
      TreeNode* target,
      int K,
      vector<int>& answer
  ) {
      if (node == nullptr) {
          return -1;
      }
61
      // Descendants exactly K levels below
      // target are collected immediately.
      if (node == target) {
          collectDown(
              node,
              K,
              answer
          );
70
          return 0;
      }
73
      int leftResult =
          findTarget(
              node->left,
              target,
              K,
              answer
          );
81
      if (leftResult != -1) {
          int currentDistance =
              leftResult + 1;
85
          // The ancestor itself is valid when
          // its distance from target equals K.
          if (currentDistance == K) {
              answer.push_back(node->val);
          } else {
              // The opposite subtree is searched for
              // nodes completing the remaining distance.
              collectDown(
                  node->right,
                  K - currentDistance - 1,
                  answer
              );
          }
99
          return currentDistance;
      }
102
      int rightResult =
          findTarget(
              node->right,
              target,
              K,
              answer
          );
110
      if (rightResult != -1) {
          int currentDistance =
              rightResult + 1;
114
          if (currentDistance == K) {
              answer.push_back(node->val);
          } else {
              // When target is on the right,
              // the left subtree becomes opposite.
              collectDown(

Complexity Analysis

Let N be the number of nodes in the binary tree and H be the height of the tree.

Time Complexity: O(N), where N is the number of nodes in the binary tree. The relevant subtrees are processed a constant number of times overall while the target is located and nodes at the required distances are collected.

Space Complexity: O(H), where H is the height of the binary tree, due to the recursion stack. This becomes O(N) for a skewed tree and O(log N) for a balanced tree.

Interview follow-up Questions

Some valid nodes may lie above the target or inside another subtree reached through an ancestor. A downward-only traversal would miss such nodes.

Print All Nodes at Distance K from a Target Node

Example 1

Example 2

Approach 1

Algorithm

Dry Run

Solution

Complexity Analysis

Approach 2

Algorithm

Dry Run

Solution

Complexity Analysis

Approach 3

Algorithm

Dry Run

Solution

Complexity Analysis

Interview follow-up Questions

Read Similar Blogs

Parent and Distance - Introduction

Check If a Binary Tree Is a Subtree of Another Tree

Count Nodes in a Complete Binary Tree

Print All Nodes at Distance K from a Target Node

Example 1

Example 2

Approach 1

Algorithm

Dry Run

Solution

Complexity Analysis

Approach 2

Algorithm

Dry Run

Solution

Complexity Analysis

Approach 3

Algorithm

Dry Run

Solution

Complexity Analysis

Interview follow-up Questions

Why can DFS not simply be started downward from the target?

Why is a visited set required in the BFS approaches?

What happens when K = 0?

Why is the opposite subtree searched in the recursive approach?

Which approach is generally preferred in interviews?

Read Similar Blogs

Parent and Distance - Introduction

Check If a Binary Tree Is a Subtree of Another Tree

Count Nodes in a Complete Binary Tree