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)

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(K == 0)
    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;
     return result;

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

      pair currentNode = q.front(); q.pop();
      if(currentNode.second == 0)
        // add left and right if they exist, and decrement the
        // distance
        Node * cur = currentNode.first;
           q.push(make_pair(cur->left, currentNode.second - 1));
           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