All Nodes Distance K in Binary Tree Interview Question for Meta

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.

You can return the answer in any order.

For example:

Consider the following binary tree:

If target is the node with value 5 and k is 2, the nodes at distance 2 from the target node are 7, 4, and 1. Therefore, the expected output is [7, 4, 1].

Another example:

If root is a single node tree with value 1, target is the node with value 1, and k is 3, the expected output is [] because there are no nodes at distance 3 from the target node.

Write a function that takes the root of a binary tree, the target node, and an integer k as input, and returns a list of the values of all nodes that are distance k from the target node.

Naive Approach: Brute Force

One straightforward way to solve this problem is to perform a breadth-first search (BFS) or depth-first search (DFS) starting from every node in the tree. For each node, calculate its distance from the target node. If the distance is equal to k, add the node's value to the result list. This is a brute-force approach because it explores all possible paths from every node.

Big O Runtime: O(N^2) where N is the number of nodes in the tree. We potentially visit every node for every other node.

Big O Space Usage: O(N) in the worst-case scenario where the tree is heavily unbalanced.

Optimal Approach: Graph Transformation + BFS

A more efficient approach involves two main steps:

Convert the Tree to a Graph: Transform the binary tree into an undirected graph. Each node in the tree becomes a node in the graph, and the edges represent parent-child relationships. This can be done using a DFS or BFS traversal of the tree.
BFS from Target: Perform a Breadth-First Search (BFS) starting from the target node in the constructed graph. Keep track of the distance from the target node as you explore. When you encounter a node at a distance of k, add its value to the result.

Edge Cases:

Empty Tree: If the tree is empty, return an empty list.
Target Not Found: The problem states that the target is in the tree, so we don't need to explicitly handle this, but it's good to keep in mind.
k = 0: If k is 0, return a list containing only the target node's value.
k > Height of Tree: If k is greater than the maximum possible distance from the target, the result list may be empty or contain only nodes that are actually that far away.

Big O Runtime: O(N), where N is the number of nodes in the tree. The graph construction is O(N) and the BFS is O(N).

Big O Space Usage: O(N). We need space to store the graph (adjacency list), the queue for BFS, and the result list.

from collections import defaultdict, deque

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def distanceK(root: TreeNode, target: TreeNode, k: int) -> list[int]:
    graph = defaultdict(list)

    def build_graph(node):
        if not node:
            return
        if node.left:
            graph[node].append(node.left)
            graph[node.left].append(node)
            build_graph(node.left)
        if node.right:
            graph[node].append(node.right)
            graph[node.right].append(node)
            build_graph(node.right)

    build_graph(root)

    q = deque([(target, 0)])
    visited = {target}
    result = []

    while q:
        node, dist = q.popleft()
        if dist == k:
            result.append(node.val)
        elif dist > k:
          break

        for neighbor in graph[node]:
            if neighbor not in visited:
                visited.add(neighbor)
                q.append((neighbor, dist + 1))

    return result

All Nodes Distance K in Binary Tree

Solution

Naive Approach: Brute Force

Optimal Approach: Graph Transformation + BFS