All Nodes Distance K in Binary Tree Interview Question for Meta

from collections import defaultdict

# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> list[int]:
        # Build the graph representation of the tree.
        graph = defaultdict(list)

        def build_graph(node, parent):
            if node:
                if parent:
                    graph[node.val].append(parent.val)
                    graph[parent.val].append(node.val)
                build_graph(node.left, node)
                build_graph(node.right, node)

        build_graph(root, None)

        # BFS to find nodes at distance k from the target.
        queue = [(target.val, 0)]
        visited = {target.val}
        result = []

        while queue:
            node, dist = queue.pop(0)
            if dist == k:
                result.append(node)
            elif dist < k:
                for neighbor in graph[node]:
                    if neighbor not in visited:
                        queue.append((neighbor, dist + 1))
                        visited.add(neighbor)

        return result

# Example Usage:
# Create the tree: [3,5,1,6,2,0,8,null,null,7,4]
root = TreeNode(3)
root.left = TreeNode(5)
root.right = TreeNode(1)
root.left.left = TreeNode(6)
root.left.right = TreeNode(2)
root.right.left = TreeNode(0)
root.right.right = TreeNode(8)
root.left.right.left = TreeNode(7)
root.left.right.right = TreeNode(4)

# Find the target node (node with value 5)
target = root.left

# Find all nodes at distance 2 from the target node
k = 2
solution = Solution()
result = solution.distanceK(root, target, k)
print(result)  # Output: [7, 4, 1]


# Naive approach (not implemented but conceptually explained)
# A naive solution might involve traversing the entire tree for each node to calculate
# the distance to the target. This would have a higher time complexity.

# Time Complexity Analysis:
# Building the graph takes O(N) time, where N is the number of nodes in the tree.
# The BFS traversal takes O(N) time in the worst case, as we might visit each node.
# Therefore, the overall time complexity is O(N).

# Space Complexity Analysis:
# The graph stores at most 2 edges for each node, so the space complexity for the graph
# is O(N).
# The queue can contain at most all the nodes in the tree, so the space complexity for the queue
# is O(N).
# The visited set can contain at most all the nodes in the tree, so the space complexity for the visited set
# is O(N).
# Therefore, the overall space complexity is O(N).

# Edge Cases:
# 1. Empty Tree: If the root is None, the function should return an empty list.
# 2. Target Not Found: If the target node is not in the tree, the function should return an empty list.
# 3. k = 0: If k is 0, the function should return a list containing only the target node's value.
# 4. k > Height of Tree: If k is greater than the height of the tree, the function may return an empty list if no nodes are found at that distance.
# 5. Tree with one node: If the tree only contains the root node, and k > 0, the function should return an empty list.

Explanation:

This code provides a solution to find all nodes at a distance k from a given target node in a binary tree. Here's a breakdown of the approach:

Build the Graph:
- The binary tree is converted into an undirected graph where each node in the tree becomes a node in the graph, and edges connect parent and child nodes. This is done using a Depth-First Search (DFS) traversal of the tree.
- A defaultdict(list) called graph is used to store the graph. The keys are node values, and the values are lists of neighboring node values.
Breadth-First Search (BFS):
- A BFS is performed starting from the target node to find all nodes at distance k. The queue stores tuples of (node_value, distance). The visited set is used to keep track of visited nodes to avoid cycles.
- The BFS explores the graph level by level. When a node at distance k is found, its value is added to the result list.
Time Complexity:
- Building the graph takes O(N) time, where N is the number of nodes in the tree.
- The BFS traversal takes O(N) time in the worst case, as we might visit each node.
- Therefore, the overall time complexity is O(N).
Space Complexity:
- The graph stores at most 2 edges for each node, so the space complexity for the graph is O(N).
- The queue can contain at most all the nodes in the tree, so the space complexity for the queue is O(N).
- The visited set can contain at most all the nodes in the tree, so the space complexity for the visited set is O(N).
- Therefore, the overall space complexity is O(N).
Edge Cases:
- Empty Tree: If the root is None, the function should return an empty list.
- Target Not Found: If the target node is not in the tree, the function should return an empty list.
- k = 0: If k is 0, the function should return a list containing only the target node's value.
- k > Height of Tree: If k is greater than the height of the tree, the function may return an empty list if no nodes are found at that distance.
- Tree with one node: If the tree only contains the root node, and k > 0, the function should return an empty list.

Example:

For the given example:

root = [3,5,1,6,2,0,8,null,null,7,4]
target = 5
k = 2

The nodes at distance 2 from node 5 are 7, 4, and 1. Therefore, the output is [7, 4, 1].