Find Leaves of Binary Tree

LeavesTraversal.java

//  Find Leaves of Binary Tree

// LeafTraversal
// Concept Problem


//Given a binary tree, find all leaves and then remove those leaves. 
// Then repeat the previous steps until the tree is empty.

public List<List<TreeNode>> leafTraversal(TreeNode root) {
  if(root == null) return root;
  List<List<TreeNode>> leaves = new ArrayList<List<TreeNode>>();
  leafTraversalRecurse(root, leaves);
  return leaves;
}

private int /*returns level*/ leafTraversalRecurse(TreeNode root, List<List<TreeNode>> leaves) {
  if(root == null) return -1; // return -1 for leaf nodes so the level is 0. 
  
  int leftLevel = leafTraversalRecurse(root.left, leaves);
  int rightLevel = leafTraversalRecurse(root.right, leaves);
  int level = Math.max(leftLevel, rightLevel) + 1;
  
  if(level == leaves.size()) { // can be replaced by (level >= leaves.size())
     leaves.add(new ArrayList<TreeNode>());
  }
  leaves.get(level).add(root);
  root.left = root.right = null; // remove the leaves
  return level;
}