863. All Nodes Distance K in Binary Tree - Detailed Explanation

Problem Statement

Given the root of a binary tree, a target node target, and an integer k, the task is to find all nodes in the tree that are distance k from the target node. The distance between two nodes is defined by the number of edges in the shortest path connecting them. You can think of this as all nodes that can be reached by taking exactly k steps from the target node. The result should be returned as a list of node values (in any order).

• Distance 0 means just the target node itself.
• Distance 1 means the immediate neighbors of the target (its parent or its direct children).
• Distance k means nodes that you reach via exactly k parent or child pointers from the target.

Examples:

1. Example 1:
   Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
   Output: [7, 4, 1]
   Explanation: Starting from node with value 5, the nodes at distance 2 are:

   • 5 → 2 → 7 (two steps: from 5 down to 2, then 2 down to 7)
   • 5 → 2 → 4 (two steps: from 5 down to 2, then 2 down to 4)
   • 5 → 3 → 1 (two steps: from 5 up to its parent 3, then 3 down to 1)
     These nodes have values 7, 4, and 1.

2. Example 2:
   Input: root = [1,2,3], target = 1, k = 1
   Output: [2, 3]
   Explanation: The target is the root node (value 1). Nodes at distance 1 are its direct children with values 2 and 3.

3. Example 3 (Edge Case):
   Input: root = [1], target = 1, k = 3
   Output: [] (empty list)
   Explanation: There is only one node in the tree. Any distance k > 0 will yield no nodes. In this case, no node exists at distance 3 from the target.

Constraints:

• The number of nodes in the tree is in the range [1, 500]. (The tree is non-empty in given cases.)
• Node values are unique integers in the range 0 <= Node.val <= 500 .
• target is guaranteed to be a value of a node present in the tree.
• 0 <= k <= 1000 . (Note: k can be larger than the height of the tree, in which case the output will be empty, as seen in Example 3.)

The output list can be returned in any order – the focus is just on which node values are at the required distance.

Hints to Guide the Thought Process

• Think of the Tree as a Graph: In a binary tree, each node has up to two children, and (except for the root) one parent. If we consider parent-child connections as undirected edges, the tree becomes an undirected graph. Finding nodes at distance k from the target is then like a graph problem: all nodes at exactly k edges away from a starting node. This hint suggests that algorithms like BFS (breadth-first search) on graphs might be applicable.

• Traversing Both Downwards and Upwards: Normally, a tree traversal (DFS or BFS) goes downward (from parent to children). Here, because we need to find nodes that could be above the target (ancestors) as well as below (descendants), we need a way to move up from a node to its parent. Consider how you might move upwards in a tree – perhaps by storing parent pointers or using recursion to backtrack.

• Use a Parent Map: One effective strategy is to perform a traversal of the tree to record each node’s parent in a hash map or dictionary. This way, given any node (including the target), you can quickly find its parent and treat that as an adjacent node in the graph sense. For example, for each node, store parent[node] = its_parent. This will allow you to move from a node to its parent in O(1) time during the search.

• Avoid Revisiting Nodes: When exploring in an undirected graph, you must be careful not to go in circles. In this tree graph, if you go from a child to its parent, you don’t want to immediately go back down to the child. Use a visited set to mark nodes you have already seen. This ensures you don’t revisit nodes and get stuck in an infinite loop between child and parent.

• Two Main Approaches – BFS or DFS: Both breadth-first search (BFS) and depth-first search (DFS) can find nodes at distance k, but in different ways. BFS will naturally find nodes by increasing distance (level by level), which fits this problem well. DFS can be used with careful accounting of distances or by splitting the problem (e.g., find target node, then explore). Think about which might be easier to implement for you:

  • BFS: starting from the target, expand outward step by step until you reach distance k.
  • DFS: you could recursively search for target and then explore in all directions for nodes at the required distance.

• Start with Simpler Cases: Imagine simpler versions:

  • What if target was the root? (Then the answer is just all nodes at depth k from the root.)
  • What if k = 0? (Then the answer is just the target node itself.)
  • These can often be handled easily and can build intuition for the general case.

By considering these hints, you can formulate a plan to solve the problem by transforming it into a graph traversal task and then deciding between BFS or DFS to collect the nodes at the desired distance.

Approach Explanation

There are two common approaches to solve this problem: using Depth-First Search (DFS) or Breadth-First Search (BFS). In both approaches, a key step is to allow traversal upwards to parent nodes, which is not directly possible in a standard tree traversal. We achieve this by first constructing a graph representation of the tree, or at least a parent reference map, so that from any node we can move to its children or its parent. After that, we either perform a BFS or DFS starting from the target node to find all nodes at distance k.

Depth-First Search (DFS) Approach

Idea: Use DFS to traverse the tree and prepare for the search, and then use DFS again to find all nodes at the required distance. This approach usually consists of two phases:

1. Parent Mapping (DFS traversal): First, do a DFS traversal of the entire tree to record each node’s parent in a map or dictionary. This can be done with a simple recursive function. For example, when visiting node curr, if it has a left child, record parent[left_child] = curr and if it has a right child, record parent[right_child] = curr. Continue this for all nodes. After this step, you have a parent map that allows you to move from any node to its parent in O(1) time.
2. DFS from Target: Now, perform a DFS starting from the target node to find all nodes at distance k. In this DFS, you need to explore in three directions from any given node: left child, right child, and parent (using the map from step 1). Use a parameter in the DFS function to keep track of the current distance from the target.

   • If the current distance is equal to k, record the node’s value as part of the result.

   • If the current distance is less than k, continue DFS deeper into all unvisited neighbors (children or parent).

   • If the current distance exceeds k or the node is None or already visited, backtrack (prune that path).

3. Mark Visited: It’s crucial to mark nodes as visited when you traverse them in this search to avoid going in circles (for example, going from a child to parent and back to child). This can be done with a set of visited nodes. Once you traverse a node, add it to the set. In the DFS recursion, if you encounter a node that’s already in the visited set, you return immediately without further processing.
4. Collect Results: Any time your DFS reaches exactly distance k (and not beyond), add that node’s value to the result list. At the end of the DFS process, the result list will contain all nodes that are exactly k away from the target.

Step-by-Step (DFS Approach):

1. Build Parent Map (DFS1): Traverse the tree with DFS to map each node to its parent. For example, start with dfs(node=root, parent=null). In the DFS:

   • If node is not null, set parent_map[node] = parent.
   • Recurse dfs(node.left, parent=node) and dfs(node.right, parent=node) to fill the map for children.

2. Initialize Structures: Prepare an empty result list result = [] and a visited set visited = set().

3. DFS from Target (DFS2): Define a recursive function dfs_from_target(node, dist) that explores from the current node which is dist edges away from the original target:

   • If node is null or already in visited, return (stop exploring this path).

   • Mark node as visited by adding it to the set.

   • If dist == k, append node.val to result (we found a node at exact distance).

   • If dist < k, call dfs_from_target on all neighbors: node.left, node.right, and parent_map[node], with dist+1.

4. Start DFS2 at Target: Call dfs_from_target(target_node, 0) to start exploring from the target (distance 0).

5. Result: After the DFS completes, result will contain all node values at distance k. Return this list.

Why it works: The first DFS builds the necessary parent links so that the second DFS can move in all directions from the target. The second DFS then exhaustively searches all paths from the target, but the visited set ensures it doesn’t repeat paths or traverse a node more than once. This effectively explores the tree as an undirected graph (each connection from parent to child can be seen as bidirectional for the purpose of finding all nodes within k steps). By stopping the search when the distance exceeds k, we limit our traversal to only relevant nodes.

Breadth-First Search (BFS) Approach

Idea: Convert the problem into a graph BFS problem. We treat the binary tree like an undirected graph by adding edges between parent and child. Then, we perform a BFS starting from the target node and go outwards level by level until we reach distance k. BFS naturally finds nodes in order of increasing distance from the start node, which fits perfectly for this problem.

Graph Construction: We need to be able to get the neighbors of any given node (its left child, right child, and parent). There are two ways to achieve this:

• Explicit Graph (Adjacency List): Create an adjacency list (a dictionary or map) where each node maps to a list of its neighbors. You can populate this by a DFS or BFS traversal of the tree:

  • For each node, if it has a left child, add an undirected edge: add left child to the node’s neighbor list and the node to the left child’s neighbor list. Do the same for the right child.

  • This yields a graph structure where each connection (parent-child in the tree) is represented twice (child as neighbor of parent, parent as neighbor of child).

• Parent Map (Implicit Graph): Alternatively, similar to the DFS approach, store parent pointers in a map. You can then get neighbors of a node by looking at its children (from its left and right pointers) and its parent (from the map). This avoids storing a full adjacency list, since the tree already gives child pointers.

Using either representation, the neighbor-finding operation will be the same: given a node, its neighbors for the BFS are [node.left, node.right, parent_map[node]] (ignoring any None or visited neighbors).

BFS from Target: Once the graph connections are ready, do the following:

1. Initialize BFS: Create a queue for BFS. Start by enqueueing the target node with distance 0. Also initialize a visited set and add the target to it (marking it visited immediately).

2. Level-by-Level Traversal: Perform BFS by dequeuing nodes. For each node dequeued with current distance d:

   • If d == k, you have reached nodes at the target distance. You can collect this node’s value (and continue to see if other nodes in the queue are also at distance k). Important: Once you reach distance k, you know that all further nodes popped from the queue (in this BFS process) will be distance k or greater. You can choose to stop the BFS here or simply skip adding further neighbors.
   • If d < k, enqueue all unvisited neighbors of this node with distance d+1. Neighbors include the left child, right child, and parent (if they exist and haven’t been visited). Mark each neighbor as visited when you enqueue it.

3. Stopping Condition: Continue the BFS until you either reach distance k or the queue is empty. A common pattern is to use a loop that runs k times, each time expanding one level of neighbors. Alternatively, you check distance in the queue:

   • You can use a counter or include the distance in the queue tuples. For example, push (node, dist) in the queue. When you pop from queue, if dist == k, record the value.
   • Another approach: perform BFS level by level. Increment a distance counter each time you finish exploring one level of the tree. When the counter reaches k, stop and collect all nodes currently in the queue (they will all be k distance away).

4. Result Collection: All nodes that have distance exactly k from the target will have been reached when d == k. Collect those node values into a result list. If the BFS finishes (queue empties) before reaching k (which means the tree does not have nodes that far away), then the result list remains empty (as in Example 3).

Step-by-Step (BFS Approach):

1. Build Parent Map: (If not building full adjacency list) Do a DFS or BFS on the tree to fill a map parent[node] = parent_node for each node (similar to DFS approach step 1).

2. Initialize Queue: Use a queue to start BFS from the target. For example, queue = deque(), enqueue target with distance 0 (you can store a tuple (target_node, 0)).

3. Initialize Visited: visited = { target_node } to mark the target as visited. This prevents immediate backtracking to it.

4. BFS Loop: Set current_distance = 0. While the queue is not empty and current_distance < k:

   • Determine the number of nodes at the current level: level_size = len(queue).

   • For each node in this level, dequeue it. Enqueue all its unvisited neighbors (children + parent) into the queue. Mark them visited.

   • After processing all nodes of this level, increment current_distance += 1. Now the queue contains all nodes at current_distance from target.

   • If current_distance == k, you can break out of the loop or proceed to collect results.

5. Collect Results: The queue now contains all nodes that are exactly k distance from target (if k levels were fully processed). Extract their values into a list. If the queue became empty before reaching k, the result is empty.

6. Return the Result List.

Why it works: BFS is an ideal fit because it explores neighbors in waves. The first wave (distance 1) explores all nodes 1 step away, the second wave (distance 2) explores all nodes 2 steps away, and so on. By the time BFS has taken k steps, it has reached all nodes at distance k. Using the parent map allows BFS to go upward from the target to its parent, effectively exploring the tree in all directions. The visited set ensures each node is processed only once, preventing infinite loops in the undirected graph formed by the tree connections.

Comparison of DFS vs BFS approaches: Both approaches require an initial traversal to record parent pointers (or build adjacency). The main difference lies in how the distance-k nodes are collected:

• DFS uses recursion or an explicit stack to depth-first search all paths up to length k. It may be slightly more complex to implement because you need to manage the distance count in recursive calls and stop exactly at k.

• BFS uses a queue to naturally traverse layer by layer. It can be a bit more straightforward to stop at exactly k by using level counts.

• Either way, the time complexity and overall work done are similar. It often comes down to which paradigm you’re more comfortable with. For clarity, many prefer BFS for this particular problem, as it aligns with the “find nodes by distance” requirement.

Code Implementation

Below are implementations in Python and Java for this problem. Both implementations follow the approach of first mapping parent pointers (to allow upward traversal), and then using BFS from the target node to collect all nodes at distance k. We include some test cases to verify the solutions.

Python Implementation