Problem Description

You are given a binary tree with a root node, a target node that exists somewhere in the tree, and an integer k. Your task is to find all nodes that are exactly k distance away from the target node.

The distance between two nodes in a tree is defined as the minimum number of edges you need to traverse to go from one node to the other. For example, a node's immediate children and parent are at distance 1, their children/parent are at distance 2, and so on.

The key insight here is that nodes at distance k from the target can be:

• In the subtree rooted at the target (going downward)
• In other parts of the tree (going upward through the parent and then possibly down other branches)

The solution uses two DFS traversals:

1. First DFS (dfs): Creates a hash table g that maps each node to its parent. This allows us to traverse upward in the tree, not just downward through children.

2. Second DFS (dfs2): Starting from the target node, explores all directions (left child, right child, and parent) to find nodes at exactly distance k. It:

   • Decrements k at each step as we move away from the target
   • When k reaches 0, we've found a node at the desired distance
   • Uses the fa (father/parent) parameter to avoid revisiting the node we came from

The algorithm effectively treats the tree as an undirected graph where we can move both up and down, finding all nodes at the specified distance from the target node. The result can be returned in any order.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: While the problem presents a binary tree, we need to treat it as a graph because we need to traverse both downward (to children) and upward (to parent nodes) from the target node. Trees are special cases of graphs.

Is it a tree?

• Yes: The problem explicitly states we're working with a binary tree structure, even though we need to traverse it like an undirected graph to find nodes at distance k.

DFS

• Yes: We arrive at DFS as our algorithm choice. This makes sense because:
  • We need to traverse the entire tree first to build parent relationships
  • We need to explore all paths from the target node to find nodes at exactly distance k
  • DFS allows us to systematically explore each path to the required depth

Conclusion: The flowchart correctly leads us to use DFS for this problem. The solution employs DFS twice:

1. First DFS traversal to map each node to its parent (building the graph representation)
2. Second DFS traversal starting from the target node to find all nodes at distance k by exploring in three directions (left child, right child, and parent)

The DFS pattern is ideal here because we need to:

• Systematically explore paths of specific length (k edges)
• Track visited nodes to avoid cycles when traversing upward through parents
• Build a complete mapping of the tree structure before searching for the target nodes

Intuition

The key insight is recognizing that finding nodes at distance k from a target requires movement in all directions - not just downward through children, but also upward through parents.

In a standard binary tree, we can only traverse downward from parent to children. However, nodes at distance k from our target could be:

• Below the target in its subtree
• Above the target (ancestors)
• In a completely different branch (reached by going up to an ancestor, then down a different path)

This limitation leads us to think: what if we could move both up and down in the tree? This transforms our tree traversal problem into a graph traversal problem where each edge is bidirectional.

To enable upward movement, we need to know each node's parent. Since the tree nodes only store references to their children, we must first traverse the entire tree to build a parent mapping. This is where the first DFS comes in - it creates a hash table g that maps each node to its parent.

Once we have parent relationships, we can treat the tree as an undirected graph. Starting from the target node, we can explore in three directions: left child, right child, and parent. This is essentially a bounded DFS where we:

• Keep track of how many steps we've taken from the target
• Stop exploring a path when we've moved exactly k steps
• Avoid revisiting nodes we came from (using the fa parameter to prevent cycles)

The beauty of this approach is that it naturally handles all cases - whether the nodes are in the target's subtree, among its ancestors, or in distant branches. By treating the tree as an undirected graph and using DFS with a distance counter, we systematically find all nodes at the exact required distance.

Learn more about Tree, Depth-First Search, Breadth-First Search and Binary Tree patterns.

Solution Approach

The implementation uses DFS with a hash table to solve this problem in two phases:

Phase 1: Build Parent Mapping

The first DFS function dfs(root, fa) traverses the entire tree to create a parent mapping:

• We maintain a hash table g where g[node] = parent_node
• Starting from the root with fa = None (root has no parent)
• For each node, we store its parent in the hash table: g[root] = fa
• Recursively process left and right children, passing the current node as their parent

This gives us the ability to traverse upward from any node in the tree.

Phase 2: Find Nodes at Distance k

The second DFS function dfs2(root, fa, k) finds all nodes at distance k from the target:

• Start from the target node with k as the remaining distance to travel
• Base cases:
  • If root is None, return (reached a null node)
  • If k == 0, we've found a node at the exact distance - add root.val to the answer
• Recursive case: explore in three directions
  • root.left: go to left child
  • root.right: go to right child
  • g[root]: go to parent node
• For each direction, we only explore if nxt != fa to avoid going back to the node we came from
• Recursively call with k - 1 since we're moving one step further from the target

Key Data Structures:

• Hash table g: Maps each node to its parent, enabling bidirectional traversal
• List ans: Collects the values of all nodes at distance k

Algorithm Flow:

1. Call dfs(root, None) to build the complete parent mapping
2. Initialize empty result list ans
3. Call dfs2(target, None, k) to find all nodes at distance k
4. Return the collected results

The time complexity is O(n) where n is the number of nodes, as we visit each node at most twice (once in each DFS). The space complexity is O(n) for the parent mapping hash table and the recursion stack.

Example Walkthrough

Let's walk through a concrete example to illustrate how this algorithm finds all nodes at distance k from a target.

Consider this binary tree:

        3
       / \
      5   1
     / \   \
    6   2   8
       / \
      7   4

Goal: Find all nodes at distance k=2 from target node 5.

Phase 1: Build Parent Mapping

Starting dfs(root=3, fa=None):

• Store g[3] = None (root has no parent)
• Recurse left: dfs(5, 3)
  • Store g[5] = 3
  • Recurse left: dfs(6, 5) → Store g[6] = 5
  • Recurse right: dfs(2, 5) → Store g[2] = 5
    • Process children: g[7] = 2, g[4] = 2
• Recurse right: dfs(1, 3)
  • Store g[1] = 3
  • Recurse right: dfs(8, 1) → Store g[8] = 1

Parent mapping complete:

g = {3: None, 5: 3, 1: 3, 6: 5, 2: 5, 7: 2, 4: 2, 8: 1}

Phase 2: Find Nodes at Distance k=2

Starting dfs2(root=5, fa=None, k=2):

Since k=2 (not 0 yet), we explore three directions from node 5:

1. Left child (node 6): dfs2(6, fa=5, k=1)

   • k=1, so continue exploring
   • Left child: None (skip)
   • Right child: None (skip)
   • Parent: g[6] = 5, but fa=5, so skip (don't go back)
   • Dead end

2. Right child (node 2): dfs2(2, fa=5, k=1)

   • k=1, so continue exploring
   • Left child (node 7): dfs2(7, fa=2, k=0)
     • k=0! Add 7 to answer
   • Right child (node 4): dfs2(4, fa=2, k=0)
     • k=0! Add 4 to answer
   • Parent: g[2] = 5, but fa=5, so skip

3. Parent (node 3): dfs2(3, fa=5, k=1)

   • k=1, so continue exploring
   • Left child: g[3.left] = 5, but fa=5, so skip (don't revisit)
   • Right child (node 1): dfs2(1, fa=3, k=0)
     • k=0! Add 1 to answer
   • Parent: g[3] = None, skip (null)

Final Result: ans = [7, 4, 1]

The algorithm successfully found all three nodes at distance 2 from node 5:

• Nodes 7 and 4 (in target's subtree, going down)
• Node 1 (in a different branch, reached by going up to parent 3, then down to sibling branch)

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

The algorithm consists of two depth-first search (DFS) traversals:

• The first DFS (dfs) visits each node exactly once to build the parent mapping in the dictionary g, taking O(n) time.
• The second DFS (dfs2) starts from the target node and explores nodes at distance k. In the worst case, it may visit all n nodes (when k is large enough to reach all nodes), taking O(n) time.

Therefore, the overall time complexity is O(n) + O(n) = O(n).

Space Complexity: O(n)

The space usage comes from:

• The dictionary g that stores parent pointers for each node, requiring O(n) space.
• The recursion call stack for DFS, which in the worst case (skewed tree) can go up to depth n, using O(n) space.
• The answer list ans that stores at most n node values, using O(n) space in the worst case.

Therefore, the overall space complexity is O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Infinite Loop Due to Missing Visited Tracking

The most critical pitfall in this problem is creating an infinite loop when traversing the tree as an undirected graph. Without proper tracking of where we came from, the algorithm can bounce back and forth between a node and its parent indefinitely.

Incorrect Implementation:

def find_nodes_at_distance_k(node, remaining_distance):
    if node is None:
        return
    if remaining_distance == 0:
        result.append(node.val)
        return
  
    # WRONG: This will cause infinite recursion between parent and child
    find_nodes_at_distance_k(node.left, remaining_distance - 1)
    find_nodes_at_distance_k(node.right, remaining_distance - 1)
    find_nodes_at_distance_k(parent_map[node], remaining_distance - 1)

Solution: Always pass and check the previous node parameter to avoid revisiting the node you came from:

def find_nodes_at_distance_k(node, previous, remaining_distance):
    # ... base cases ...
  
    for next_node in (node.left, node.right, parent_map[node]):
        if next_node != previous:  # Critical check to prevent cycles
            find_nodes_at_distance_k(next_node, node, remaining_distance - 1)

2. Forgetting to Handle None Parent in Parent Map

When exploring from a node, accessing parent_map[node] for the root node returns None, which is correct. However, if not handled properly in the traversal logic, this can cause issues.

Potential Issue:

# If you try to access properties of parent_map[root] without None check
parent = parent_map[node]
if parent.val == something:  # This will crash for root node

Solution: The current implementation handles this correctly by treating None as a valid value in the loop and letting the base case handle it:

for next_node in (node.left, node.right, parent_map[node]):
    if next_node != previous:  # None is handled naturally here
        find_nodes_at_distance_k(next_node, node, remaining_distance - 1)

3. Not Building Complete Parent Map Before Search

A common mistake is trying to build the parent map while searching for nodes at distance k, or only building it partially.

Incorrect Approach:

def distanceK(self, root, target, k):
    # WRONG: Trying to build parent map only from target
    parent_map = {}
    build_parent_map(target, None)  # This only maps nodes in target's subtree
    # ...

Solution: Always build the complete parent map starting from the root before beginning the search:

# Correct: Build complete parent map from root
build_parent_map(root, None)
# Then search from target
find_nodes_at_distance_k(target, None, k)

4. Modifying Distance Parameter Incorrectly

Some implementations might try to increment distance instead of decrementing it, or forget to modify it at all.

Incorrect:

# WRONG: Incrementing instead of decrementing
find_nodes_at_distance_k(next_node, node, remaining_distance + 1)

# WRONG: Forgetting to change distance
find_nodes_at_distance_k(next_node, node, remaining_distance)

Solution: Always decrement the remaining distance as you move away from the target:

find_nodes_at_distance_k(next_node, node, remaining_distance - 1)

5. Using a Set for Visited Nodes Instead of Previous Node

While using a visited set works, it's less efficient and uses more memory than simply tracking the previous node.

Less Efficient Approach:

def find_nodes_at_distance_k(node, visited, remaining_distance):
    if node is None or node in visited:
        return
    visited.add(node)
    # ... rest of logic

Better Solution: The current implementation elegantly uses just the previous parameter to prevent cycles, which is more memory-efficient and cleaner.