Find Nodes Distance K

Tags

Find Nodes Distance K

Category

Binary Trees

Difficulty

Hard

Problem Statement

Given a binary tree, a target node within the tree, and a non-negative integer k, find all nodes that are exactly distance k from the target node. The distance between two nodes is the number of edges on the path connecting them. Return the values of all such nodes.

Intuition

In a tree, distance k nodes can be in two directions from the target: downward in the target's subtree, or upward through the target's ancestors. Finding nodes below the target is straightforward DFS. Finding nodes above requires the ability to traverse from child to parent, which a standard binary tree does not support. The solution is to first create parent pointers (via a DFS/BFS pass), then perform a BFS or DFS from the target node in all three directions (left child, right child, parent) for exactly k steps.

Approach

  1. Traverse the entire tree to build a mapping from each node to its parent node.
  2. Starting from the target node, perform a BFS (or DFS with depth tracking).
  3. Use a visited set to avoid revisiting nodes (since we can now traverse in both directions).
  4. At each BFS level, expand to all unvisited neighbors: left child, right child, and parent.
  5. When the BFS reaches depth k, all nodes at that level are the answer.
  6. Return the values of those nodes.

Pseudocode

function findNodesDistanceK(root, target, k):
    parents = {}
    populateParents(root, null, parents)

    visited = set()
    queue = [(target, 0)]
    visited.add(target)

    while queue is not empty:
        node, distance = queue.dequeue()

        if distance == k:
            // All remaining nodes in queue at this distance are results
            result = [node.value]
            for each remaining (n, d) in queue where d == k:
                result.append(n.value)
            return result

        for neighbor in [node.left, node.right, parents[node]]:
            if neighbor is not null and neighbor not in visited:
                visited.add(neighbor)
                queue.enqueue((neighbor, distance + 1))

    return []

function populateParents(node, parent, parents):
    if node is null:
        return
    parents[node] = parent
    populateParents(node.left, node, parents)
    populateParents(node.right, node, parents)

Time & Space Complexity

  • Time: O(n) where n is the number of nodes. Building the parent map takes O(n), and the BFS visits each node at most once.
  • Space: O(n) for the parent map, the visited set, and the BFS queue.

Key Insights

  • The parent map effectively converts the tree into an undirected graph, enabling traversal in all directions.
  • BFS is the natural choice because it explores nodes level by level (i.e., by distance), making it trivial to stop at exactly distance k.
  • An alternative approach avoids parent pointers by using a recursive method that computes distances through the return values, but BFS with parent pointers is more intuitive.
  • The visited set is essential to prevent infinite loops when traversing the parent edge back to a child.

Edge Cases

  • k = 0 — the only result is the target node itself.
  • The target is the root — no parent traversal is possible; all results are in the subtree.
  • The target is a leaf — all results are through the parent direction.
  • k is larger than the tree's diameter — return an empty list.
  • The tree has only one node — if k = 0, return that node; otherwise return an empty list.

Linked mentions