All Nodes Distance K in Binary Tree

medium tree dfs bfs hash map

Problem

Given the root of a binary tree, a target node, and integer k, return all node values that are at distance exactly k from target (counting parent edges as well).

Inputtree = [3,5,1,6,2,0,8,null,null,7,4], target=5, k=2

Output[7,4,1]

From node 5: distance 2 reaches 7 (via 2), 4 (via 2), and 1 (via 3).

level-order (use 'null')

target value

from collections import deque
def distanceK(root, target, k):
    parent = {}
    def dfs(node, par):
        if not node: return
        parent[node] = par
        dfs(node.left, node); dfs(node.right, node)
    dfs(root, None)
    seen = {target}
    q = deque([target])
    d = 0
    while q and d < k:
        for _ in range(len(q)):
            cur = q.popleft()
            for nb in (cur.left, cur.right, parent[cur]):
                if nb and nb not in seen:
                    seen.add(nb); q.append(nb)
        d += 1
    return [n.val for n in q]

function distanceK(root, target, k) {
  const parent = new Map();
  (function dfs(node, par) {
    if (!node) return;
    parent.set(node, par);
    dfs(node.left, node); dfs(node.right, node);
  })(root, null);
  const seen = new Set([target]);
  let q = [target], d = 0;
  while (q.length && d < k) {
    const next = [];
    for (const cur of q) {
      for (const nb of [cur.left, cur.right, parent.get(cur)]) {
        if (nb && !seen.has(nb)) { seen.add(nb); next.push(nb); }
      }
    }
    q = next; d++;
  }
  return q.map(n => n.val);
}

class Solution {
    Map<TreeNode, TreeNode> parent = new HashMap<>();
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        dfs(root, null);
        Set<TreeNode> seen = new HashSet<>(); seen.add(target);
        Deque<TreeNode> q = new ArrayDeque<>(); q.offer(target);
        int d = 0;
        while (!q.isEmpty() && d < k) {
            int sz = q.size();
            for (int i = 0; i < sz; i++) {
                TreeNode cur = q.poll();
                for (TreeNode nb : new TreeNode[]{ cur.left, cur.right, parent.get(cur) }) {
                    if (nb != null && seen.add(nb)) q.offer(nb);
                }
            }
            d++;
        }
        List<Integer> res = new ArrayList<>();
        for (TreeNode n : q) res.add(n.val);
        return res;
    }
    void dfs(TreeNode n, TreeNode p) { if (n == null) return; parent.put(n, p); dfs(n.left, n); dfs(n.right, n); }
}

vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
    unordered_map<TreeNode*, TreeNode*> parent;
    function<void(TreeNode*, TreeNode*)> dfs = [&](TreeNode* n, TreeNode* p) {
        if (!n) return; parent[n] = p; dfs(n->left, n); dfs(n->right, n);
    };
    dfs(root, nullptr);
    unordered_set<TreeNode*> seen{target};
    queue<TreeNode*> q; q.push(target);
    int d = 0;
    while (!q.empty() && d < k) {
        int sz = (int)q.size();
        for (int i = 0; i < sz; i++) {
            TreeNode* cur = q.front(); q.pop();
            for (TreeNode* nb : { cur->left, cur->right, parent[cur] }) {
                if (nb && !seen.count(nb)) { seen.insert(nb); q.push(nb); }
            }
        }
        d++;
    }
    vector<int> res;
    while (!q.empty()) { res.push_back(q.front()->val); q.pop(); }
    return res;
}

Explanation

Distance in a tree usually only flows downward (parent to child), but here distance can also go up through parents. The fix is to turn the tree into a regular graph where every node knows its neighbors in all three directions.

First, a quick DFS builds a parent map so every node remembers who its parent is. Now from any node we can reach three neighbors: left, right, and parent.

Then we do a breadth-first search starting at the target. BFS expands outward one ring at a time, so after k rounds the queue holds exactly the nodes whose distance is k.

A seen set stops us from walking back to nodes we already visited, which is what keeps the rings clean and prevents infinite loops.

Example: target 5, k = 2. Distance 1 reaches 6, 2, and 3 (its parent). Distance 2 reaches 7, 4 (children of 2) and 1 (other child of 3), giving [7,4,1].

Time: O(n) Space: O(n)