All Nodes Distance K in Binary Tree

mediumAmazon-tagged: Trees & Matrix (mock day)112 stepsJava · Python · C++

Return every node exactly k edges from a target node. Tree edges only point down, so first record each node's parent. Then BFS from the target in three directions (left, right, parent) for k levels.

▶ Watch the animated trace

Traced input: root = [3, 5, 1, 6, 2, 0, 8, null, null, 7, 4], target = 5, k = 2 → [1, 4, 7]. Original problem: leetcode.com/problems/all-nodes-distance-k-in-binary-tree/

Complexity

BestO(n)

AverageO(n)

WorstO(n)

SpaceO(n)

One pass for parents, one BFS.

Key idea

Adding parent links turns the tree into an undirected graph.
BFS level k from the target is exactly the answer.
A visited set stops the BFS from walking back the way it came.

Common mistakes

Only searching down from the target misses nodes reached through the parent.
Forgetting the visited set loops between a child and its parent.

Solution

Java

List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
    Map<TreeNode, TreeNode> parent = new HashMap<>();
    Deque<TreeNode> stack = new ArrayDeque<>();
    stack.push(root);
    while (!stack.isEmpty()) {
        TreeNode n = stack.pop();
        if (n.left != null) { parent.put(n.left, n); stack.push(n.left); }
        if (n.right != null) { parent.put(n.right, n); stack.push(n.right); }
    }
    Queue<TreeNode> queue = new ArrayDeque<>();
    Set<TreeNode> seen = new HashSet<>();
    queue.offer(target);
    seen.add(target);
    for (int d = 0; d < k && !queue.isEmpty(); d++) {
        for (int size = queue.size(); size > 0; size--) {
            TreeNode n = queue.poll();
            for (TreeNode next : new TreeNode[]{n.left, n.right, parent.get(n)}) {
                if (next != null && seen.add(next)) queue.offer(next);
            }
        }
    }
    List<Integer> res = new ArrayList<>();
    for (TreeNode n : queue) res.add(n.val);
    return res;
}

Python

def distanceK(root: TreeNode, target: TreeNode, k: int) -> List[int]:
    parent = {}
    stack = [root]
    while stack:
        n = stack.pop()
        for child in (n.left, n.right):
            if child: parent[child] = n; stack.append(child)
    queue, seen = deque([target]), {target}
    for _ in range(k):
        if not queue: break
        for _ in range(len(queue)):
            n = queue.popleft()
            for nxt in (n.left, n.right, parent.get(n)):
                if nxt and nxt not in seen: seen.add(nxt); queue.append(nxt)
    return [n.val for n in queue]

C++

vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
    unordered_map<TreeNode*, TreeNode*> parent;
    vector<TreeNode*> stack{root};
    while (!stack.empty()) {
        TreeNode* n = stack.back(); stack.pop_back();
        if (n->left) { parent[n->left] = n; stack.push_back(n->left); }
        if (n->right) { parent[n->right] = n; stack.push_back(n->right); }
    }
    queue<TreeNode*> q;
    unordered_set<TreeNode*> seen{target};
    q.push(target);
    for (int d = 0; d < k && !q.empty(); d++) {
        for (int size = q.size(); size > 0; size--) {
            TreeNode* n = q.front(); q.pop();
            for (TreeNode* next : {n->left, n->right, parent[n]}) {
                if (next && seen.insert(next).second) q.push(next);
            }
        }
    }
    vector<int> res;
    for (; !q.empty(); q.pop()) res.push_back(q.front()->val);
    return res;
}

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
k = 0 The target itself.	`root = [1, 2, 3], target = 2, k = 0`	`[2]`
k larger than the tree The BFS runs out of nodes.	`root = [1], target = 1, k = 3`	`[]`
Answer through the parent 2 is a leaf; its only neighbour is its parent.	`root = [0, 1, null, 3, 2], target = 2, k = 1`	`[1]`