Find all nodes in a binary tree that are at the distance K from the root

Input: A binary tree, number K
Output: A list of all nodes that are at the distance K from the root
Problem: Find all nodes in a binary tree that are at the distance K from the root
Complexity: Time O(n), Memory 0(n)
Example:

     
L:0       A
...    /     \
L:1   B       C
...   \     / \
L:2    D   E   F
...   / \
L:3  G   H


For K = 2, we got [D, E, F], L:N stands for the level N.

Firstly, let’s assume our tree node is defined in the following way:

struct Node
{
    struct Node * left;
    struct Node * right;
    int value;
};

Basically what we are looking for is the level K, as all nodes on this level will be at the distance of K from the root. In case of this problem there is really not much to think. The only choice we need to make is if go for breadth-first-search, or depth-first-search algorithm. Then we go to the distance K from our root node and we add the nodes at this level to the result. Below you can find both of the algorithms implemented.

DFS approach:

void findKFarFromRoot(Node * node, int K, vector result)
{
    if(!node)return;
    if(K == 0)
       result.append(node);
    else {
       findKFarFromRoot(node->left, K - 1, result);
       findKFarFromRoot(node->right, K - 1, result);
    }
}

BFS approach:

vector findKFarFromRoot(Node * root, int K)
{
   vector<pair<Node *, int> > result;
   if(!root)
     return result;

   queue q;
   q.push(make_pair(root, K));

   while(!q.empty())
   {
      pair currentNode = q.front(); q.pop();
      if(currentNode.second == 0)
        result.push_back(currentNode);
      else
      {
        // add left and right if they exist, and decrement the
        // distance
        Node * cur = currentNode.first;
        if(cur->left)
           q.push(make_pair(cur->left, currentNode.second - 1));
        if(cur->right)
           q.push(make_pair(cur->right, currentNode.second - 1));
      }
   }
   return result
}

Obviously, both algorithms has complexity O(n) as they visit every node once. Please free to comment if you see mistakes, or something is not clear. But keep in mind it’s a rather c++ based pseudo code that real solution (althought should be correct). Both solutions works for a more generic case: they return all the nodes at the distance K from a node given as a parameter.

Best Regards

Joe Kidd

Advertisements