All Nodes Distance K in Binary Tree

Given the root of a binary tree, a target node, and an integer k, return all the nodes that are distance k from the target node.

Example 1

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

Example 2

Input: root=[1], target=1, k=3 → Output: []

Solution Approach

Convert the tree to an undirected graph (tracking parent pointers). Then BFS from target node with distance k.

DFS to build parent map: parent[node] = its parent.
BFS from target node, tracking visited nodes.
Expand each level; after k levels, collect all nodes in the queue.

Java

import java.util.*;

class Solution {
    Map<TreeNode, TreeNode> parent = new HashMap<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        buildParent(root, null);
        Queue<TreeNode> q = new LinkedList<>();
        Set<TreeNode> visited = new HashSet<>();
        q.add(target); visited.add(target);
        int dist = 0;
        while (!q.isEmpty()) {
            if (dist == k) {
                List<Integer> res = new ArrayList<>();
                for (TreeNode n : q) res.add(n.val);
                return res;
            }
            int size = q.size();
            for (int i = 0; i < size; i++) {
                TreeNode n = q.poll();
                for (TreeNode nb : new TreeNode[]{n.left, n.right, parent.get(n)}) {
                    if (nb != null && !visited.contains(nb)) { q.add(nb); visited.add(nb); }
                }
            }
            dist++;
        }
        return new ArrayList<>();
    }
    private void buildParent(TreeNode node, TreeNode par) {
        if (node == null) return;
        parent.put(node, par);
        buildParent(node.left, node); buildParent(node.right, node);
    }
}

Time Complexity: O(N)
Space Complexity: O(N)