LeetCode 366. Find Leaves of Binary Tree

原题网址：https://leetcode.com/problems/find-leaves-of-binary-tree/

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

Example:
Given binary tree 

          1
         / \
        2   3
       / \     
      4   5    

Returns [4, 5, 3], [2], [1].

Explanation:

1. Remove the leaves [4, 5, 3] from the tree

          1
         / 
        2          

2. Remove the leaf [2] from the tree

          1          

3. Remove the leaf [1] from the tree

          []         

Returns [4, 5, 3], [2], [1].

方法：先检查最大深度，再根据节点所处的深度放进结果。

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    private int depth(TreeNode node) {
        if (node == null) return 0;
        int depth = 1;
        int left = depth(node.left);
        int right = depth(node.right);
        depth += Math.max(left, right);
        return depth;
    }
    
    private int depth(TreeNode node, List<List<Integer>> results) {
        if (node == null) return 0;
        int depth = 1;
        int left = depth(node.left, results);
        int right = depth(node.right, results);
        depth += Math.max(left, right);
        results.get(depth - 1).add(node.val);
        return depth;
    }
    
    public List<List<Integer>> findLeaves(TreeNode root) {
        int depth = depth(root);
        List<List<Integer>> results = new ArrayList<>(depth);
        for(int i = 0; i < depth; i++) {
            results.add(new ArrayList<>());
        }
        depth(root, results);
        return results;
    }
}

简化：

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    private int find(TreeNode node, List<List<Integer>> results) {
        if (node == null) return 0;
        int depth = Math.max(find(node.left, results), find(node.right, results));
        if (depth == results.size()) results.add(new ArrayList<>());
        results.get(depth).add(node.val);
        return depth + 1;
    }
    public List<List<Integer>> findLeaves(TreeNode root) {
        List<List<Integer>> results = new ArrayList<>();
        if (root != null) find(root, results);
        return results;
    }
}