366. Find Leaves of Binary Tree

Given a binary tree, collect a tree's nodes as if you were doing this: Collect and remove all leaves, repeat until the tree is empty.

 

Example:

Input: [1,2,3,4,5]
  
          1
         / \
        2   3
       / \     
      4   5    

Output: [[4,5,3],[2],[1]]

 

Explanation:

1. Removing the leaves [4,5,3] would result in this tree:

          1
         / 
        2          

 

2. Now removing the leaf [2] would result in this tree:

          1          

 

3. Now removing the leaf [1] would result in the empty tree:

          []          

[[3,5,4],[2],[1]], [[3,4,5],[2],[1]], etc, are also consider correct answers since per each level it doesn't matter the order on which elements are returned. 

My answer:

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71

/** * Definition for a binary tree node. * public class TreeNode { * int val; * TreeNode left; * TreeNode right; * TreeNode() {} * TreeNode(int val) { this.val = val; } * TreeNode(int val, TreeNode left, TreeNode right) { * this.val = val; * this.left = left; * this.right = right; * } * } */ class Solution { public List<List<Integer>> findLeaves(TreeNode root) { // Solution 1 // List<List<Integer>> result = new ArrayList<List<Integer>>(); // TreeNode resultNode = root; // while(resultNode != null) { // List<Integer> singleResult = new ArrayList<Integer>(); // resultNode = removeLeaves(root, singleResult); // result.add(singleResult); // } // return result; // Solution 2: add leaf to its correspondig node's height // node's height is number of edges to its deepest leaf List<List<Integer>> result = new ArrayList<List<Integer>>(); height(root, result); return result; } private int height(TreeNode node, List<List<Integer>> result) { if (node == null) { return -1; } int currentHeight = 1 + Math.max(height(node.left, result), height(node.right, result)); if (result.size() == currentHeight) { // found a new level/height for result // because largest index of result is size - 1 result.add(new ArrayList<Integer>()); } result.get(currentHeight).add(node.val); // left child is already null in last recursion, thus no need to do that // node.left = null; // node.right = null; node = null; return currentHeight; } private TreeNode removeLeaves(TreeNode node, List<Integer> leaves) { if (node == null) { return null; } // found leaf node if (node.left == null && node.right == null ) { leaves.add(node.val); return null; } TreeNode left = removeLeaves(node.left, leaves); TreeNode right = removeLeaves(node.right, leaves); node.left = left; node.right = right; return node; } }