All Nodes at Distance K in a Binary Tree

Given a binary tree with unique values, a target value in the tree, and an int k, return the values of all nodes that are exactly k edges away from the target. Paths can go up through parents as well as down.

LeetCode: 863. All Nodes Distance K in Binary Tree

Example

graph TB
    a3["3"] --> a5["5 (target)"]
    a3 --> a1["1"]
    a5 --> a6["6"]
    a5 --> a2["2"]
    a1 --> a0["0"]
    a1 --> a8["8"]
    a2 --> a7["7"]
    a2 --> a4["4"]

Text

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

7 and 4 are two edges below 5; 1 is reached by going up to 3 and down again. Find Nodes at k Distance From the Root only goes down; here the parent links matter too.

Straightforward: build a graph, then BFS

Tree nodes are objects with left and right references, and a missing child is null. The tree has no parent links, so turn it into an undirected graph (an adjacency list keyed by value). Then run a breadth-first search from the target, k levels out.

Java

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class Main {
    static class TreeNode {
        int val;
        TreeNode left, right;

        TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }

    static TreeNode node(int val, TreeNode left, TreeNode right) {
        return new TreeNode(val, left, right);
    }

    static TreeNode leaf(int val) {
        return new TreeNode(val, null, null);
    }

    static void addEdges(TreeNode n, Map<Integer, List<Integer>> graph) {
        for (TreeNode child : new TreeNode[]{n.left, n.right}) {
            if (child == null) continue;
            graph.computeIfAbsent(n.val, key -> new ArrayList<>()).add(child.val);
            graph.computeIfAbsent(child.val, key -> new ArrayList<>()).add(n.val);
            addEdges(child, graph);
        }
    }

    static List<Integer> distanceK(TreeNode root, int target, int k) {
        Map<Integer, List<Integer>> graph = new HashMap<>();
        if (root != null) {
            addEdges(root, graph);
        }
        Set<Integer> seen = new HashSet<>(List.of(target));
        List<Integer> frontier = new ArrayList<>(List.of(target));
        for (int step = 0; step < k; step++) {
            List<Integer> next = new ArrayList<>();
            for (int v : frontier) {
                for (int w : graph.getOrDefault(v, List.of())) {
                    if (seen.add(w)) {
                        next.add(w);
                    }
                }
            }
            frontier = next;
        }
        return frontier;
    }

    public static void main(String[] args) {
        TreeNode root = node(3, node(5, leaf(6), node(2, leaf(7), leaf(4))), node(1, leaf(0), leaf(8)));
        List<Integer> result = distanceK(root, 5, 2);
        Collections.sort(result);
        System.out.println(result);
    }
}

Output

[1, 4, 7]

Time: O(n). Space: O(n) for the graph.

seen.add(w) returns false for a vertex already visited, so the filter keeps the search from walking back towards the target.

Optimal: one DFS, no graph

Nodes at distance k are either below the target (depth k in its subtree), or reached by going up d edges to an ancestor and then down k - d edges into the ancestor's other subtree.

A DFS that returns the distance from each node to the target (if the target is below it) finds both kinds:

At the target: collect the nodes k levels down, and return OptionalInt.of(0).
At an ancestor that is d edges above the target (one child returned d - 1): if d == k, the ancestor itself is at distance k. If d < k, collect the nodes k - d - 1 levels down in the other child. Return OptionalInt.of(d).

Java

import java.util.ArrayList;
import java.util.List;
import java.util.OptionalInt;

public class Main {
    static class TreeNode {
        int val;
        TreeNode left, right;

        TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }

    static TreeNode node(int val, TreeNode left, TreeNode right) {
        return new TreeNode(val, left, right);
    }

    static TreeNode leaf(int val) {
        return new TreeNode(val, null, null);
    }

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

    // Returns the distance from root down to the target, if the target is in this subtree.
    static OptionalInt search(TreeNode root, int target, int k, List<Integer> out) {
        if (root == null) return OptionalInt.empty();
        if (root.val == target) {
            collectDown(root, k, out);
            return OptionalInt.of(0);
        }
        TreeNode[][] sides = {{root.left, root.right}, {root.right, root.left}}; // {side, other}
        for (TreeNode[] pair : sides) {
            OptionalInt below = search(pair[0], target, k, out);
            if (below.isPresent()) {
                int d = below.getAsInt() + 1;
                if (d == k) {
                    out.add(root.val);
                } else if (d < k) {
                    collectDown(pair[1], k - d - 1, out);
                }
                return OptionalInt.of(d);
            }
        }
        return OptionalInt.empty();
    }

    static List<Integer> distanceK(TreeNode root, int target, int k) {
        List<Integer> out = new ArrayList<>();
        search(root, target, k, out);
        return out;
    }

    public static void main(String[] args) {
        TreeNode root = node(3, node(5, leaf(6), node(2, leaf(7), leaf(4))), node(1, leaf(0), leaf(8)));
        System.out.println(distanceK(root, 5, 2));
        System.out.println(distanceK(root, 7, 3));
        System.out.println(distanceK(root, 3, 0) + " " + distanceK(root, 6, 9));
    }
}

Output

[7, 4, 1]
[6, 3]
[3] []

Time: O(n): the search visits each node once, and the collectDown calls cover disjoint subtrees. Space: O(h) for the recursion.

Edge cases

k = 0: only the target itself.
k larger than any distance: an empty result.
Target is the root: only the "below" case applies.