medium All Nodes Distance K in Binary Tree

Problem Statement

Given the root of a binary tree, a target node, and a distance k, find all the nodes in the tree that are exactly k edges away from the target node. You may return these nodes' values in any order.

Examples

Example 1

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

• Expected Output: [6, 1]
• Justification: Nodes 1 and 6 are one edge away from node 3.

Example 2

• Input: root = [1, 2, 3, null, 4, 5, 6], tarrget = 2, k = 2

• Expected Output: [3]
• Justification: Node 3 is two edges away from node 2.

Example 3

• Input: root = [1, 2, 3, 4, 5, null, null, null, null, 6, 7], target = 4, k = 3

• Expected Output: [6, 7, 3]
• Justification: Nodes 6, 7 and 3 are three edges away from node 4.

Constraints:

• The number of nodes in the tree is in the range [1, 500].
• 0 <= Node.val <= 500
• All the values Node.val are unique.
• target is the value of one of the nodes in the tree.
• 0 <= k <= 1000

Solution

To solve this problem, we need to traverse the binary tree and find all nodes that are k edges away from a given target node. We can achieve this by using a two-step approach. First, we need to convert the tree into a graph structure so that we can easily navigate through the nodes, regardless of their direction (parent or child). This graph will allow us to perform a breadth-first search (BFS) starting from the target node to find all nodes that are exactly k edges away. The BFS approach ensures that we efficiently explore all possible paths to the nodes at the desired distance, making it an effective solution for this problem.

Step-by-Step Algorithm

1. Initialize Data Structures:

   • Create a dictionary parentMap to store the parent of each node.
   • Create a list result to store the nodes at distance k.

2. Build Parent Map:

   • Define a helper method buildParentMap that takes a node and its parent as arguments.
   • For each node, store its parent in parentMap.
   • Recursively call buildParentMap for the left and right children of the current node.

3. Breadth-First Search (BFS):

   • Initialize a queue queue and add the target node to it.
   • Initialize a set seen and add the target node to it to track visited nodes.
   • Initialize a variable distance to zero.

4. Perform BFS:

   • While the queue is not empty:
     • If distance equals k, collect all nodes in the queue into result and return result.
     • For each node in the current level:
       • Dequeue the node.
       • For each neighbor (left child, right child, parent):
         • If the neighbor exists and is not in seen, add it to seen and enqueue it.
     • Increment distance by one after processing all nodes at the current distance.

5. Return Result:

   • If no nodes are found at distance k, return an empty list.

Algorithm Walkthrough

Build the Parent Map in DFS Manner

We use a DFS traversal to build the parent map, starting from the root node. Each step includes setting the parent of a node, then recursively moving to the left and right children.

1. Start DFS with Root Node (1):

   • Set the parent of Node 1 to null (no parent for root).
   • Recur on the left child (Node 2).

2. Node 2:

   • Set the parent of Node 2 to Node 1.

   • Recur on the left child (Node 4).

   • Node 4:

     • Set the parent of Node 4 to Node 2.
     • Node 4 has no children, so backtrack to Node 2.

   • Back to Node 2:

     • Now recur on the right child (Node 5).

   • Node 5:

     • Set the parent of Node 5 to Node 2.

     • Recur on the left child (Node 6).

     • Node 6:

       • Set the parent of Node 6 to Node 5.
       • Node 6 has no children, so backtrack to Node 5.

     • Back to Node 5:

       • Now recur on the right child (Node 7).

     • Node 7:

       • Set the parent of Node 7 to Node 5.
       • Node 7 has no children, so backtrack.

   • Backtrack fully to Node 1.

3. Back to Root Node (1):

   • Now recur on the right child (Node 3).

4. Node 3:

   • Set the parent of Node 3 to Node 1.
   • Node 3 has no children, so DFS is complete.

• Final Parent Map: After completing the DFS traversal, the parentMap is as follows:

{
  1: null,
  2: 1,
  3: 1,
  4: 2,
  5: 2,
  6: 5,
  7: 5
}

Queue Operations with Set Updates (BFS Traversal)

Next, we will show the state of the queue and seen set at each BFS level. We start BFS from the target node (4) and progress level-by-level until reaching the required distance K = 3.

1. Level 0 (distance = 0):

   • Queue: [4]
   • Set: {4}
   • Target node 4 has no children but has a parent (2). Add 2 to the queue and seen.
   • Increment distance to 1.

2. Level 1 (distance = 1):

   • Queue: [2]
   • Set: {2, 4}
   • Node 2 has:
     • Left child 4 (already in seen).
     • Right child 5. Add 5 to the queue and seen.
     • Parent 1. Add 1 to the queue and seen.
   • Increment distance to 2.

3. Level 2 (distance = 2):

   • Queue: [5, 1]
   • Set: {1, 2, 4, 5}
   • Process node 5:
     • Left child 6. Add 6 to the queue and seen.
     • Right child 7. Add 7 to the queue and seen.
     • Parent 2 (already in seen).
   • Process node 1:
     • Left child 2 (already in seen).
     • Right child 3. Add 3 to the queue and seen.
   • Increment distance to 3.

4. Level 3 (Distance = K(3)):

   • Queue: [6, 7, 3]
   • Set: {1, 2, 3, 4, 5, 6, 7}
   • Since we’ve reached distance = K, all nodes currently in the queue (6, 7, and 3) are at distance K = 3 from the target node 4.
   • Add nodes in the queue to the result list.

The BFS traversal ends here, and the result list [6, 7, 3] is returned as the final output.

import java.util.*;

// class TreeNode {
//     int val;
//     TreeNode left;
//     TreeNode right;
//     TreeNode(int x) { val = x; }
// }

class Solution {

  public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
    Map<TreeNode, TreeNode> parentMap = new HashMap<>();
    List<Integer> result = new ArrayList<>();

    // Populate the parent map
    buildParentMap(root, null, parentMap);

    // Perform BFS from the target node
    Queue<TreeNode> queue = new LinkedList<>();
    queue.add(target);
    Set<TreeNode> seen = new HashSet<>();
    seen.add(target);
    int distance = 0;

    while (!queue.isEmpty()) {
      if (distance == K) {
        // Collect all nodes at distance K
        for (TreeNode node : queue) {
          result.add(node.val);
        }
        return result;
      }
      int size = queue.size();
      for (int i = 0; i < size; i++) {
        TreeNode node = queue.poll();
        // Explore the left child
        if (node.left != null && seen.add(node.left)) {
          queue.add(node.left);
        }
        // Explore the right child
        if (node.right != null && seen.add(node.right)) {
          queue.add(node.right);
        }
        // Explore the parent node
        TreeNode parent = parentMap.get(node);
        if (parent != null && seen.add(parent)) {
          queue.add(parent);
        }
      }
      distance++;
    }
    return result;
  }

  // Helper method to build parent map
  private void buildParentMap(
    TreeNode node,
    TreeNode parent,
    Map<TreeNode, TreeNode> parentMap
  ) {
    if (node == null) return;
    parentMap.put(node, parent);
    buildParentMap(node.left, node, parentMap);
    buildParentMap(node.right, node, parentMap);
  }

  public static void main(String[] args) {
    Solution solution = new Solution();
    // Example 1
    TreeNode root1 = new TreeNode(1);
    root1.left = new TreeNode(2);
    root1.right = new TreeNode(3);
    root1.left.left = new TreeNode(4);
    root1.left.right = new TreeNode(5);
    root1.right.right = new TreeNode(6);
    root1.left.left.left = new TreeNode(7);
    root1.left.left.right = new TreeNode(8);
    TreeNode target1 = root1.right; // Node with value 3
    System.out.println(solution.distanceK(root1, target1, 1)); // Expected output: [1, 6]

    // Example 2
    TreeNode root2 = new TreeNode(1);
    root2.left = new TreeNode(2);
    root2.right = new TreeNode(3);
    root2.left.right = new TreeNode(4);
    root2.right.left = new TreeNode(5);
    root2.right.right = new TreeNode(6);
    TreeNode target2 = root2.left; // Node with value 2
    System.out.println(solution.distanceK(root2, target2, 2)); // Expected output: [3]

    // Example 3
    TreeNode root3 = new TreeNode(1);
    root3.left = new TreeNode(2);
    root3.right = new TreeNode(3);
    root3.left.left = new TreeNode(4);
    root3.left.right = new TreeNode(5);
    root3.left.right.left = new TreeNode(6);
    root3.left.right.right = new TreeNode(7);
    TreeNode target3 = root3.left.left; // Node with value 4
    System.out.println(solution.distanceK(root3, target3, 3)); // Expected output: [6, 7, 3]
  }
}

Complexity Analysis

Time Complexity

1. Building the Parent Map:

   • We traverse the entire tree once to build the parent map.
   • This takes time, where ( N ) is the number of nodes in the tree.

2. Breadth-First Search (BFS):

   • In the worst case, we might visit all the nodes.
   • This also takes time.

Combining both steps, the total time complexity is .

Space Complexity

1. Parent Map:

   • The parent map stores a reference for each node, resulting in space.

2. Queue and Seen Set:

   • In the worst case, the queue and seen set can store all nodes, resulting in space.

Combining both, the total space complexity is .