algoadvance

We are given the root of a binary tree and a target node, and an integer K. We need to return a list of the values of all nodes that have a distance K from the target node. The distance between two nodes is the number of edges on the shortest path connecting them.

Example

Given the tree:

      3
     / \
    5   1
   / \ / \
  6  2 0  8
    / \
   7   4

Target node = 5 and K = 2

Output: [7, 4, 1]

Clarifying Questions

  1. Can the target node have a distance of 0? Yes, if K is 0, we return the target node itself.

  2. Are all values in the tree unique? Yes, the values of the nodes in the tree are unique.

  3. What should be returned if K is larger than the height of the tree? Return an empty list as there will be no nodes at that distance.

  4. Can the root be null? Yes, if the root is null, we should return an empty list.

  5. What if the target node is not found in the tree? You can assume the input provides a valid target node present in the tree.

Strategy

  1. Convert the Tree into a Graph: First, convert the binary tree into an undirected graph representation. This allows us to naturally handle the distances using BFS.

  2. BFS Traversal: Use Breadth-First Search (BFS) starting from the target node to find all nodes that are at distance K.

  3. Edge Cases:

    • Handle when K is 0 by returning the target node as the only element in the list.
    • If the root is null, return an empty list immediately.

Code

from collections import defaultdict, deque
from typing import List, Optional

# 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: Optional[TreeNode], target: TreeNode, K: int) -> List[int]:
    if not root:
        return []
    
    # Build graph from tree
    graph = defaultdict(list)
    
    def build_graph(node, parent=None):
        if node and parent:
            graph[node.val].append(parent.val)
            graph[parent.val].append(node.val)
        if node.left:
            build_graph(node.left, node)
        if node.right:
            build_graph(node.right, node)
    
    build_graph(root)
    
    # Perform BFS from the target node
    queue = deque([(target.val, 0)])
    seen = set([target.val])
    result = []
    
    while queue:
        node, dist = queue.popleft()
        if dist == K:
            result.append(node)
        elif dist < K:
            for neighbor in graph[node]:
                if neighbor not in seen:
                    seen.add(neighbor)
                    queue.append((neighbor, dist + 1))
    
    return result

Time Complexity

Building the graph: Each node is visited once, and the edges are created for each connection. This takes O(N) time, where N is the number of nodes in the tree.

BFS Traversal: In the worst case, we visit every node and every edge, which takes O(N) time.

Overall, the time complexity is O(N).

Space Complexity: The space required for storing the graph takes O(N) space. Additionally, the BFS queue and seen set will also take up O(N) space in the worst case. Thus, the overall space complexity is O(N).

Try our interview co-pilot at AlgoAdvance.com