Print all nodes that are at distance k from a leaf node

Problem

Given a Binary Tree and a positive integer k, print all nodes that are distance k from a leaf node.

Here the meaning of distance is different from previous post. Here k distance from a leaf means k levels higher than a leaf node. For example if k is more than height of Binary Tree, then nothing should be printed. Expected time complexity is O(n) where n is the number nodes in the given Binary Tree.




Example
(Please ignore the empty node, and consider it null)

k = 1, Answer = 2, 19 , 21
k = 2, Answer = 5, 18 , 19

Solution

The idea is to traverse the tree. Keep storing all ancestors till we hit a leaf node. When we reach a leaf node, we print the ancestor at distance k. We also need to keep track of nodes that are already printed as output. For that we use a boolean array visited[].

// This function prints all nodes that are distance k from a leaf node
//   path[] - Store ancestors of a node
//   visited[] - Stores true if a node is printed as output.  A node may be k
//                 distance away from many leaves, we want to print it once 
void kDistantFromLeafUtil(Node node, int path[], bool visited[],
                          int pathLen, int k)
{
    // Base case
    if (node==null) return;
 
    // append this Node to the path array 
    path[pathLen] = node.data;
    visited[pathLen] = false;
    pathLen++;
 
    // it's a leaf, so print the ancestor at distance k only
    // if the ancestor is not already printed  
    if (node.left == null && node.right == null &&
        pathLen-k-1 >= 0 && visited[pathLen-k-1] == false)
    {
        System.out.print(path[pathLen-k-1] + " ");
        visited[pathLen-k-1] = true;
        return;
    }
 
    // If not leaf node, recur for left and right subtrees 
    kDistantFromLeafUtil(node.left, path, visited, pathLen, k);
    kDistantFromLeafUtil(node.right, path, visited, pathLen, k);
}
 
// Given a binary tree and a nuber k, print all nodes that are k
//   distant from a leaf
void printKDistantfromLeaf(Node node, int k)
{
    int[] path = new int[MAX_HEIGHT];
    boolean[] visited = new boolean[MAX_HEIGHT];
    //all the elements false in visited
    Arrays.fill(visited, false);
    kDistantFromLeafUtil(node, path, visited, 0, k);
}

References

• http://www.geeksforgeeks.org/print-nodes-distance-k-leaf-node/