863. All Nodes Distance K in Binary Tree

Traverse the binary tree twice. We convert the tree into a graph and then traverse the graph twice. The first traversal creates parent nodes, and then the second BFS goes through K layers to return all neighbors with distance k for each node.

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        def dfs(node, parent):
            if not node:
                return
 
            node.parent = parent
            dfs(node.left, node)
            dfs(node.right, node)
 
        dfs(root, None)
        queue = deque([target])
        seen = {target}
        distance = 0
 
        while queue and distance < k:
            curr_length = len(queue)
            for _ in range(curr_length):
                node = queue.popleft()
                for neighbor in [node.left, node.right, node.parent]:
                    if neighbor and neighbor not in seen:
                        seen.add(neighbor)
                        queue.append(neighbor)
            distance += 1
        
        return [node.val for node in queue]

Without mutation, create a graph structure to store the parents for each node.

class Solution:
    def distanceK(self, root: Optional[TreeNode], target: TreeNode, k: int) -> List[int]:
        # Step 1: Build adjacency list
        graph = defaultdict(list)
 
        def dfs(node, parent):
            if not node:
                return
            if parent:
                graph[node].append(parent)
                graph[parent].append(node)
            dfs(node.left, node)
            dfs(node.right, node)
 
        dfs(root, None)
 
        # Step 2: BFS from target
        queue = deque([target])
        seen = {target}
        distance = 0
 
        while queue and distance < k:
            for _ in range(len(queue)):
                node = queue.popleft()
                for neighbor in graph[node]:
                    if neighbor not in seen:
                        seen.add(neighbor)
                        queue.append(neighbor)
            distance += 1
 
        return [node.val for node in queue]