863. All Nodes Distance K in Binary Tree

Problem¶

Given a binary tree, a target node, and an integer k, return the values of every node exactly k edges from the target, in any order.

root = [3,5,1,6,2,0,8,null,null,7,4], target = node(5), k = 2
distance 2 from node(5) -> [7,4,1]

Key trick¶

Treat the tree as an undirected graph by recording each node's parent, then run BFS from the target through its left child, right child, and parent.

Trap¶

Tree nodes normally point only downward, so ignoring parent links misses nodes outside the target's subtree.
A visited set is required because parent links create cycles.
target is a node reference, not merely its value.
Stop expanding nodes once distance k is reached.

Why is it interesting?¶

It tests whether you can temporarily reinterpret a tree as a graph and apply level-order BFS.

Python solution¶

from collections import deque


class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> list[int]:
        # Build parent links so every tree edge can be traversed both ways.
        parent = {root: None}
        stack = [root]

        while stack:
            node = stack.pop()

            if node.left:
                parent[node.left] = node
                stack.append(node.left)

            if node.right:
                parent[node.right] = node
                stack.append(node.right)

        # Each BFS layer contains nodes at the same distance from target.
        q = deque([target])
        visited = {target}
        dist = 0

        while q and dist < k:
            for _ in range(len(q)):
                node = q.popleft()

                for nxt in (node.left, node.right, parent[node]):
                    if nxt and nxt not in visited:
                        visited.add(nxt)
                        q.append(nxt)

            dist += 1

        return [node.val for node in q]

Time complexity: \(O(n)\).

Space complexity: \(O(n)\).

Comment on my solution¶

Your parent-link idea and BFS choice are correct, but splitting the search into downward and upward traversals introduces several bugs:

This condition checks the wrong child:

if node.right and node.left not in visited:

It should check whether node.right was visited. In the given example, node.left is node 5, so node 1 is incorrectly skipped.

This continue prevents exploring side subtrees of an ancestor:

if node_parent and node_parent not in visited:
    visited.add(node_parent)
    q.append((node_parent, parent[node_parent], dist + 1))
    continue

When the target is the root, both BFS traversals search its descendants and may return duplicates.
When k == 0, both traversals append the target.
A single BFS over children and parent links handles every direction uniformly and avoids these special cases.

# Ideas:
# - find target
#   - while finding it, keep its parent to be able to traverse tree
#     via its parent
# - start bfs from target, k layer, both visiting children and parent

from collections import deque

# Wrong Answer: 23/57 testcases passed
# [3,5,1,6,2,0,8,null,null,7,4], 5, 2
# Output: [7,4]
# Expected: [7,4,1]
class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> list[int]:
        parent = {root: None}

        # Find target in tree and filling parent
        stack = [root]
        while target not in parent:
            node = stack.pop()
            if node.left:
                parent[node.left] = node
                stack.append(node.left)
            if node.right:
                parent[node.right] = node
                stack.append(node.right)

        res = []

        # Find nodes in target's children
        q = deque([(target, 0)]) # (node, distance from target)
        while q:
            node, dist = q.popleft()
            if dist == k:
                res.append(node.val)
                continue
            if node.left:
                q.append((node.left, dist + 1))
            if node.right:
                q.append((node.right, dist + 1))

        # Find nodes in target's ancestor and the other part of the tree
        q = deque([(target, parent[target], 0)]) # (node, parent, distance from target)
        visited = {target}
        while q:
            node, node_parent, dist = q.popleft()

            if dist == k:
                res.append(node.val)
                continue

            if node_parent and node_parent not in visited:
                visited.add(node_parent)
                q.append((node_parent, parent[node_parent], dist + 1))
                continue

            if node.left and node.left not in visited:
                visited.add(node.left)
                q.append((node.left, None, dist + 1))
            if node.right and node.left not in visited:
                visited.add(node.right)
                q.append((node.right, None, dist + 1))

        return res