Print and remove leaf nodes of given Binary Tree on each iteration

Given a binary tree, the task is to:

• Print all the leaf nodes and then delete them all. 
• Repeat this process till the tree becomes empty.

Examples:

Input:
              1
           /.    \
        2         3
    /      \   
 4          5      
Output: 
Iteration 1: 4 5 3
Iteration 2: 2
Iteration 3: 1
Explanation: The leaf nodes initially were 4, 5 and 3.
When those were deleted only node with no child is 2.
So the leaf node is 2. After 2 is deleted 1 is the only node of the tree.
So 1 is deleted in the 3rd and last iteration.

Input: 
           1
         / 
       2
     /
   3               

Output: 
Iteration 1: 3
Iteration 2: 2
Iteration 3: 1 

Naive Approach: This problem can be solved by using any traversal algorithm multiple times and at each iteration simply print and remove all the leaf nodes. Reference for this approach can be found here. 

Time Complexity: O(N²)
Auxiliary Space: O(1)

Efficient Approach: The idea to solve the problem efficiently is as follows:

Instead of doing the traversal multiple time simply traverse only once and at the same time calculate the depth of each node. 
Then sort the nodes as per depth and will get first leaf nodes at depth 0, second iteration leaf nodes at depth 1 and so on.

Follow the steps to solve the problem:

• Traverse the tree recursively using DFS.
• For any node: 
  • Find the left Subtree Depth. 
  • Find the right subtree depth.  
  • Depth of node = max(leftSubtreeDepth, rightSubtreeDepth) +1
• Store the node value in a map as per depth (depth being key and value of the node as an element in the vector)

Below is the implementation of the above approach.

C++

// C++ program to print and
// remove leaf nodes

#include <bits/stdc++.h>
using namespace std;

// A Binary Tree Node
struct Node {
    int data;
    struct Node *left, *right;
};

map<int, vector<int> > resultMap;

// Function to store depth of each nodes
int fillMap(Node* root)
{
    if (root == nullptr)
        return 0;
    int LsubTreeDepth = fillMap(root->left);
    int RsubTreeDepth = fillMap(root->right);
    int depth = max(LsubTreeDepth,
                    RsubTreeDepth)
                + 1;
    resultMap[depth].push_back(root->data);
    return depth;
}

// Print and remove leaf nodes
void printLeafNodes(Node* root)
{

    // If node is null, return
    if (!root)
        return;

    // maxDepth from tree
    int maxDepth = fillMap(root);
    for (int i = 1; i <= maxDepth; i++) {
        vector<int> tempVector
            = resultMap[i];
        int vSize = tempVector.size();

        // Print leaf nodes
        // from each iteration
        cout << "Iteration " << i << " : ";
        for (int j = 0; j < vSize; j++)
            cout << tempVector[j] << " ";
        cout << "\n";
    }
}

// Utility function to
// create a new tree node
Node* newNode(int data)
{
    Node* temp = new Node;
    temp->data = data;
    temp->left = temp->right = NULL;
    return temp;
}

// Driver Code
int main()
{
    // Create binary tree
    Node* root = newNode(1);
    root->left = newNode(2);
    root->right = newNode(3);
    root->left->left = newNode(4);
    root->left->right = newNode(5);

    // Print leaf nodes of the given tree
    printLeafNodes(root);
    return 0;
}

Java

// Java program to print and remove leaf nodes
import java.io.*;
import java.util.*;

class Node {
  int data;
  Node left, right;

  Node(int data)
  {
    this.data = data;
    left = right = null;
  }
}

class GFG {

  Map<Integer, Vector<Integer> > resultMap
    = new HashMap<>();

  // Function to store depth of each nodes
  int fillMap(Node root)
  {
    if (root == null) {
      return 0;
    }
    int LsubTreeDepth = fillMap(root.left);
    int RsubTreeDepth = fillMap(root.right);
    int depth
      = Math.max(LsubTreeDepth, RsubTreeDepth) + 1;
    resultMap
      .computeIfAbsent(depth, k -> new Vector<>())
      .add(root.data);
    return depth;
  }

  // Print and remove leaf nodes
  void printLeafNodes(Node root)
  {
    // If node is null, return
    if (root == null) {
      return;
    }

    // maxDepth from tree
    int maxDepth = fillMap(root);
    for (int i = 1; i <= maxDepth; i++) {
      Vector<Integer> tempVector = resultMap.get(i);
      int vSize = tempVector.size();

      // Print leaf nodes
      // from each iteration
      System.out.print("Iteration " + i + " : ");
      for (int j = 0; j < vSize; j++) {
        System.out.print(tempVector.get(j) + " ");
      }
      System.out.println();
    }
  }

  // Utility function to create a new tree node
  Node newNode(int data)
  {
    Node temp = new Node(data);
    temp.left = temp.right = null;
    return temp;
  }

  public static void main(String[] args)
  {
    GFG tree = new GFG();

    // Create binary tree
    Node root = tree.newNode(1);
    root.left = tree.newNode(2);
    root.right = tree.newNode(3);
    root.left.left = tree.newNode(4);
    root.left.right = tree.newNode(5);

    // Print leaf nodes of the given tree
    tree.printLeafNodes(root);
  }
}

// This code is contributed by lokesh.

Python

class Node:
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None

result_map = {}

# Function to store depth of each nodes
def fill_map(root):
    if root is None:
        return 0
    LsubTreeDepth = fill_map(root.left)
    RsubTreeDepth = fill_map(root.right)
    depth = max(LsubTreeDepth, RsubTreeDepth) + 1
    if depth in result_map:
        result_map[depth].append(root.data)
    else:
        result_map[depth] = [root.data]
    return depth

# Print and remove leaf nodes
def print_leaf_nodes(root):

    # If node is null, return
    if not root:
        return

    # maxDepth from tree
    max_depth = fill_map(root)
    for i in range(1, max_depth+1):
        temp_vector = result_map[i]
        v_size = len(temp_vector)

        # Print leaf nodes
        # from each iteration
        print("Iteration ", i)
        for j in range(v_size):
            print(temp_vector[j])
        print()

# Utility function to create a new tree node
def new_node(data):
    temp = Node(data)
    temp.left = None
    temp.right = None
    return temp

# Driver code
if __name__ == '__main__':

    # Create binary tree
    root = new_node(1)
    root.left = new_node(2)
    root.right = new_node(3)
    root.left.left = new_node(4)
    root.left.right = new_node(5)

    # Print leaf nodes of the given tree
    print_leaf_nodes(root)

    # This code is contributed by aadityamaharshi21.

C#

// C# program to implement above approach
using System;
using System.Collections;
using System.Collections.Generic;

class GFG
{

    // A Binary Tree Node
    public class Node {
        public int data;
        public Node left, right;
    };

    static Dictionary<int, List<int>> resultMap = new Dictionary<int, List<int>>();

    // Function to store depth of each nodes
    static int fillMap(Node root)
    {
        if (root == null)
            return 0;

        int LsubTreeDepth = fillMap(root.left);
        int RsubTreeDepth = fillMap(root.right);
        int depth = Math.Max(LsubTreeDepth,  RsubTreeDepth) + 1;

        if(!resultMap.ContainsKey(depth)){
            resultMap.Add(depth, new List<int>());
        }
        resultMap[depth].Add(root.data);
        return depth;
    }

    // Print and remove leaf nodes
    static void printLeafNodes(Node root)
    {

        // If node is null, return
        if (root == null)
            return;

        // maxDepth from tree
        int maxDepth = fillMap(root);
        for (int i = 1; i <= maxDepth; i++) {
            List<int> tempVector = resultMap[i];
            int vSize = tempVector.Count;

            // Print leaf nodes
            // from each iteration
            Console.Write("Iteration " + i + " : ");
            for (int j = 0 ; j < vSize ; j++)
                Console.Write(tempVector[j] + " ");
            Console.WriteLine("");
        }
    }

    // Utility function to
    // create a new tree node
    static Node newNode(int data)
    {
        Node temp = new Node();
        temp.data = data;
        temp.left = temp.right = null;
        return temp;
    }

    public static void Main(string[] args){

        // Create binary tree
        Node root = newNode(1);
        root.left = newNode(2);
        root.right = newNode(3);
        root.left.left = newNode(4);
        root.left.right = newNode(5);

        // Print leaf nodes of the given tree
        printLeafNodes(root);

    }
}

// This code is contributed by entertain2022.

Javascript

       // JavaScript program to print and
       // remove leaf nodes
       class Node {
           constructor(data) {
               this.data = data;
               this.left = null;
               this.right = null;
           }
       }

       const resultMap = new Map();

       // Function to store depth of each nodes
       function fillMap(root) {
           if (!root) return 0;
           const LsubTreeDepth = fillMap(root.left);
           const RsubTreeDepth = fillMap(root.right);
           const depth = Math.max(LsubTreeDepth, RsubTreeDepth) + 1;
           if (!resultMap.has(depth)) {
               resultMap.set(depth, []);
           }
           resultMap.get(depth).push(root.data);
           return depth;
       }

       // Print and remove leaf nodes
       function printLeafNodes(root)
       {

           // If node is null, return
           if (!root) return;

           // maxDepth from tree
           const maxDepth = fillMap(root);
           for (let i = 1; i <= maxDepth; i++) {
               const tempArray = resultMap.get(i);
               console.log(`Iteration ${i} : ${tempArray.join(' ')}<br>`);
           }
       }

       // Utility function to create a new tree node
       function newNode(data) {
           return new Node(data);
       }

       // Create binary tree
       const root = newNode(1);
       root.left = newNode(2);
       root.right = newNode(3);
       root.left.left = newNode(4);
       root.left.right = newNode(5);

       // Print leaf nodes of the given tree
       printLeafNodes(root);

// This code is contributed by Potta Lokesh

Output

Iteration 1 : 4 5 3 
Iteration 2 : 2 
Iteration 3 : 1 

Time Complexity: O(N)
Auxiliary Space: O(N)