Problem Statement in English

You’re given the root of a binary tree, the value of a target node target, and an integer k.

Return an array of the values of all nodes that have a distance k from the target node, in any order.

Approach

If you look at the tree as graph, you’ll see that it’s easier to just traverse outwards from the target node. It’s quite easy to keep track of distances, and the implementation is just a matter of doing a DFS or BFS.

And that’s it.

Solution in Python


class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        if not k:
            return [target.val]
            
        adjList = defaultdict(list)

        def dfs(node):
            if node.left:
                adjList[node.val].append(node.left.val)
                adjList[node.left.val].append(node.val)
                dfs(node.left)

            if node.right:
                adjList[node.val].append(node.right.val)
                adjList[node.right.val].append(node.val)
                dfs(node.right)

        dfs(root)

        stack = [(target.val, 0)]
        res = []
        seen = set()

        while stack:
            v, d = stack.pop()
            seen.add(v)
            
            for nei in adjList[v]:
                if nei not in seen:
                    if d + 1 == k:
                        res.append(nei)
                    else:
                        stack.append((nei, d + 1))

        return res

Complexity

  • Time: $O(n)$
    Since we visit each node once.

  • Space: $O(n)$
    For the adjacency list and the stack.

And we are done.