863

863. All Nodes Distance K in Binary Tree

We’re asked to return the values of the nodes that are distance k from a given target node. We can accomplish this by conducting a DFS through the binary tree to assign each node a parent node. This allows us to treat the binary tree as an undirected graph.

With this updated graph, we can condcut a BFS from the target node outwards, conducting k steps. Once we’ve done k steps, we return the values of the nodes still remaining in the queue - these are the nodes at the level that are k distance away from the target.

The solution is as follows:

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        stack = [(root, None)]
 
        while stack:
            node, parent = stack.pop()
            node.parent = parent
            if node.left:
                stack.append((node.left, node))
            if node.right:
                stack.append((node.right, node))
 
        distance = 0
        seen = {target}
        queue = [target]
 
        while queue and distance < k:
            curr_queue, queue = queue, []
 
            for node in curr_queue:
                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]

_ Time Complexity:

O(n) - DFS and BFS both take O(n) time.

_ Space Complexity:

O(n) - We maintain a stack, queue, and seen set that all take up O(n) space.