All Nodes Distance K in Binary Tree – C++ Solution

Problem Description

So, you’ve got a binary tree, and you’re trying to find all the nodes that are exactly K distance away from a given target node. Sounds simple, right? Well, it’s like trying to find your friend in a crowded mall when you only have a vague idea of where they might be. You know they’re somewhere in the vicinity, but good luck figuring out which store they’re hiding in!

Imagine you’re at a family reunion, and you want to find all your relatives who are sitting exactly 2 tables away from you. You can’t just yell out their names; you need a strategy! This problem is essentially the same, but instead of tables, we have nodes, and instead of relatives, we have… well, nodes again.

Code Solution


class Solution {
 public:
  vector distanceK(TreeNode* root, TreeNode* target, int k) {
    vector ans;
    unordered_map nodeToDist;  // {node: distance to target}

    getDists(root, target, nodeToDist);
    dfs(root, k, 0, nodeToDist, ans);
    return ans;
  }

 private:
  void getDists(TreeNode* root, TreeNode* target,
                unordered_map& nodeToDist) {
    if (root == nullptr)
      return;
    if (root == target) {
      nodeToDist[root] = 0;
      return;
    }

    getDists(root->left, target, nodeToDist);
    if (const auto it = nodeToDist.find(root->left); it != nodeToDist.cend()) {
      // The target is in the left subtree.
      nodeToDist[root] = it->second + 1;
      return;
    }

    getDists(root->right, target, nodeToDist);
    if (const auto it = nodeToDist.find(root->right); it != nodeToDist.cend())
      // The target is in the right subtree.
      nodeToDist[root] = it->second + 1;
  }

  void dfs(TreeNode* root, int k, int dist,
           unordered_map& nodeToDist, vector& ans) {
    if (root == nullptr)
      return;
    if (const auto it = nodeToDist.find(root); it != nodeToDist.cend())
      dist = it->second;
    if (dist == k)
      ans.push_back(root->val);

    dfs(root->left, k, dist + 1, nodeToDist, ans);
    dfs(root->right, k, dist + 1, nodeToDist, ans);
  }
};

Approach Explanation

The approach taken in this solution is two-fold. First, we calculate the distance of each node from the target node using a depth-first search (DFS). This is done in the getDists function, which populates a map with nodes and their respective distances to the target.

Next, we perform another DFS to find all nodes that are exactly K distance away from the target. This is done in the dfs function, where we check if the current node’s distance matches K and, if so, we add it to our answer list.

Time and Space Complexity

Time Complexity: O(N), where N is the number of nodes in the binary tree. We traverse each node a couple of times.

Space Complexity: O(N) for storing the distances in the unordered map and the recursion stack.

Real-World Example

Think of a social network where you want to find all your friends who are two degrees away from you. You know your direct friends (distance 1), and you want to find their friends (distance 2). This problem is similar; you start from a target node (yourself) and look for all nodes (friends) that are exactly K connections away.