Topics › Binary Trees › Hard Problems

Print all nodes at a distance of K in BT

Hard TUF ID 121 A2Z

LeetCode ↗ Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

You are given the root of a binary tree, a target node inside it, and an integer k. Return the values of all nodes whose distance from target is exactly k, where distance is the number of edges on the path between two nodes. The path may go up through parents as well as down through children. Any order is accepted; the Java code returns the values sorted so that results are easy to compare.

Constraints (LeetCode): up to 500 nodes, all values distinct, 0 <= k <= 1000, and target is in the tree.

Example 1

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

target = 5, k = 2   ->  [1, 4, 7]

7 and 4 are two steps below 5. 1 is two steps away going up to 3 and then down to the right.

Example 2

Same tree, target = 5, k = 3. The answer is [0, 8]: up to 3, down to 1, then down once more. Nothing below 5 is three steps deep.

Approach 1: Brute force, distances from root paths

Intuition. The distance between two nodes in a tree goes through their lowest common ancestor. If we know each node’s path from the root, the lowest common ancestor of two nodes is the last node their paths share. With common shared nodes, the distance between nodes whose paths have lengths a and b (counted in nodes) is a + b - 2 * common: both depths are measured from the root, and the shared part is counted twice too many.

So we record the root path of every node in one DFS, then test each node against the target’s path.

Algorithm.

1. DFS from the root with a shared path; store a copy of the path for every node.
2. Let pt be the target’s path.
3. For every node with path pv, count the shared prefix common. If |pt| + |pv| - 2 * common == k, add the node.
4. Sort and return.

static List<Integer> bruteForce(TreeNode root, TreeNode target, int k) {
    Map<TreeNode, List<TreeNode>> pathOf = new HashMap<>();
    collectPaths(root, new ArrayList<>(), pathOf);
    List<Integer> ans = new ArrayList<>();
    List<TreeNode> pt = pathOf.get(target);
    if (pt == null) return ans;                               // target is not in the tree
    for (Map.Entry<TreeNode, List<TreeNode>> e : pathOf.entrySet()) {
        List<TreeNode> pv = e.getValue();
        int common = 0;                                       // nodes shared by both root paths
        while (common < pt.size() && common < pv.size() && pt.get(common) == pv.get(common)) common++;
        if (pt.size() + pv.size() - 2 * common == k) ans.add(e.getKey().val);
    }
    Collections.sort(ans);
    return ans;
}

static void collectPaths(TreeNode node, List<TreeNode> path, Map<TreeNode, List<TreeNode>> pathOf) {
    if (node == null) return;
    path.add(node);
    pathOf.put(node, new ArrayList<>(path));
    collectPaths(node.left, path, pathOf);
    collectPaths(node.right, path, pathOf);
    path.remove(path.size() - 1);
}

Complexity. Time O(n * h): each node’s path has up to h entries, both when copying and when comparing. Space O(n * h) for the stored paths. Fine for 500 nodes, wasteful in general.

Approach 2: Parent pointers and a BFS from the target

Intuition. The real obstacle is that tree edges only point downward, while distance ignores direction. Fix that by recording each node’s parent. Now every node has up to three neighbours (left, right, parent), and the tree is just an undirected graph. Nodes at distance exactly k in a graph are what a BFS from target reaches in round k.

A BFS processed level by level does exactly that. Run k rounds; the queue then holds precisely the nodes at distance k. A visited set is required: without it, the walk would go from a node to its parent and straight back to the node.

Algorithm.

1. DFS once to fill parent[node].
2. Put target in the queue and in visited.
3. Repeat k times (stopping early if the queue empties): pop the current level, and push every unvisited neighbour among left, right, parent.
4. The nodes left in the queue are the answer.

static List<Integer> parentMapBfs(TreeNode root, TreeNode target, int k) {
    Map<TreeNode, TreeNode> parent = new HashMap<>();
    fillParents(root, null, parent);
    List<Integer> ans = new ArrayList<>();
    if (!parent.containsKey(target)) return ans;
    Set<TreeNode> visited = new HashSet<>();
    Queue<TreeNode> q = new ArrayDeque<>();
    q.add(target);
    visited.add(target);
    int dist = 0;
    while (!q.isEmpty() && dist < k) {                        // after the loop, q holds exactly the ring at distance k
        int size = q.size();
        for (int i = 0; i < size; i++) {
            TreeNode node = q.poll();
            for (TreeNode next : new TreeNode[]{node.left, node.right, parent.get(node)}) {
                if (next != null && visited.add(next)) q.add(next);   // add() is false for an already visited node
            }
        }
        dist++;
    }
    for (TreeNode node : q) ans.add(node.val);
    Collections.sort(ans);
    return ans;
}

static void fillParents(TreeNode node, TreeNode par, Map<TreeNode, TreeNode> parent) {
    if (node == null) return;
    parent.put(node, par);
    fillParents(node.left, node, parent);
    fillParents(node.right, node, parent);
}

Dry run on Example 1 (target = 5, k = 2):

After two rounds the queue is [7, 4, 1], which sorts to [1, 4, 7].

Complexity. Time O(n): building the map and the BFS each touch every node at most once. Space O(n) for the map, the set, and the queue.

Approach 3: Optimal space, one DFS without a parent map

Intuition. We can avoid storing parents by letting the recursion do the climbing. Every node on the path from the root to target is an ancestor at some distance d from it. For such an ancestor A:

• if d == k, A itself is in the answer;
• if d < k, the answer also contains nodes in A’s other subtree (the one not containing target). Stepping from A into that child costs one edge, so we need nodes k - d - 1 levels below that child.

Plus the nodes exactly k levels below target itself. That covers everything: any node is either below target, or its path to target turns at some ancestor A and enters A’s other subtree, or it is such an A.

To know d, the recursive function returns the distance from the current node down to target, or -1 if target is not below it. The child that returns a non-negative distance tells us which side holds target, so we collect from the other side.

Algorithm.

1. distanceToTarget(node): return -1 for null.
2. If node == target, collect nodes k levels down from it and return 0.
3. Recurse left. If it returned l >= 0, set d = l + 1. Add node if d == k, or collect k - d - 1 levels below node.right if d < k. Return d.
4. Otherwise do the same with the right child and node.left.
5. Return -1 if target was found in neither subtree.

static List<Integer> optimal(TreeNode root, TreeNode target, int k) {
    List<Integer> ans = new ArrayList<>();
    distanceToTarget(root, target, k, ans);
    Collections.sort(ans);
    return ans;
}

/** Returns the distance from node down to target, or -1 when target is not in node's subtree. */
static int distanceToTarget(TreeNode node, TreeNode target, int k, List<Integer> ans) {
    if (node == null) return -1;
    if (node == target) {
        collectDown(node, k, ans);                            // nodes k levels below the target
        return 0;
    }
    int left = distanceToTarget(node.left, target, k, ans);
    if (left >= 0) {
        int d = left + 1;                                     // distance from node to target
        if (d == k) ans.add(node.val);
        else if (d < k) collectDown(node.right, k - d - 1, ans);   // one more edge to step into the other side
        return d;
    }
    int right = distanceToTarget(node.right, target, k, ans);
    if (right >= 0) {
        int d = right + 1;
        if (d == k) ans.add(node.val);
        else if (d < k) collectDown(node.left, k - d - 1, ans);
        return d;
    }
    return -1;
}

static void collectDown(TreeNode node, int depth, List<Integer> ans) {
    if (node == null || depth < 0) return;
    if (depth == 0) { ans.add(node.val); return; }
    collectDown(node.left, depth - 1, ans);
    collectDown(node.right, depth - 1, ans);
}

Dry run on Example 1 (target = 5, k = 2):

Node 3 returns 1 to nobody (it is the root), and the sorted answer is [1, 4, 7].

Complexity. Time O(n): distanceToTarget visits each node at most once (it never searches below target, nor the right subtree once target was found on the left), and collectDown visits each node at most once because it only enters subtrees that do not contain target. Space O(h) for the recursion, with no map or set.

Edge cases and pitfalls

• k = 0. The answer is [target]. In the BFS no rounds run; in the DFS collectDown(target, 0) adds it.
• k larger than any distance. The answer is empty. The BFS must stop when the queue empties, not only after k rounds.
• Forgetting visited. Without it, the BFS returns to nodes it came from and reports wrong distances.
• Collecting on the wrong side. In Approach 3, collecting from the subtree that contains target would count those nodes with the wrong distance. Always collect from the other child.
• target is the root or a leaf. Both are in the tests. From a leaf everything is reached by going up first.

Related problems

• Minimum time taken to burn the BT from a given Node runs the same BFS until the whole tree is covered.
• LCA in BT is the lowest common ancestor behind the distance formula.
• Print root to leaf path in BT builds the root paths used by the brute force.
• Level Order Traversal is the level-by-level BFS used in Approach 2.