LeetCode Solutions Blog

All Nodes Distance K in Binary Tree

Number: 893

Difficulty: Medium

Paid? No

Problem Description

Given the root of a binary tree, a target node (identified by its value), and an integer k, return an array of values of all nodes that have a distance k from the target node. The answer can be returned in any order.

Key Insights

Convert the binary tree into an undirected graph so that we can easily traverse both down to the children and up to the parent.
Use a DFS or simple recursive traversal to create an adjacency list representing the graph.
Perform a BFS starting from the target node to explore nodes level-by-level until reaching distance k.
Use a visited set (or equivalent) to avoid revisiting nodes and to prevent cycles.

Space and Time Complexity

Time Complexity: O(N), where N is the number of nodes in the tree. Each node is visited once when constructing the graph and once in the BFS. Space Complexity: O(N), for storing the graph and the additional data structures used in BFS and DFS.

Solution

The solution involves two main steps:

Build an undirected graph representation of the binary tree using a DFS. For each node, add the node’s value as a key in an adjacency list and connect it with its children (if any). This creates bidirectional links between parent and child.
Use a BFS starting from the target node and explore nodes level-by-level. Keep a counter for the current distance, and once the counter reaches k, return all node values found at that level. To avoid processing nodes more than once, maintain a set of visited nodes.

Code Solutions

# Definition for a binary tree node.
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: int, k: int):
    from collections import defaultdict, deque
    
    # Build an adjacency list to represent the graph
    graph = defaultdict(list)
    
    # Helper function to build the graph using DFS
    def buildGraph(node, parent):
        if node is None:
            return
        if parent is not None:
            # Add bidirectional connection between node and parent
            graph[node.val].append(parent.val)
            graph[parent.val].append(node.val)
        buildGraph(node.left, node)
        buildGraph(node.right, node)
    
    buildGraph(root, None)
    
    # BFS from target node
    visited = set([target])
    queue = deque([target])
    current_distance = 0
    
    while queue:
        if current_distance == k:
            # Return all nodes at distance k
            return list(queue)
        size = len(queue)
        for _ in range(size):
            current = queue.popleft()
            for neighbor in graph[current]:
                if neighbor not in visited:
                    visited.add(neighbor)
                    queue.append(neighbor)
        current_distance += 1
        
    return []  # If no nodes at distance k

# Example usage:
# Construct tree or use a tree builder before calling distanceK function.