650 · 二叉树叶子顺序遍历

650 · 二叉树叶子顺序遍历

描述
给定一个二叉树，像这样收集树节点：收集并移除所有叶子，重复，直到树为空。

样例
样例1
输入： {1,2,3,4,5}
输出： [[4, 5, 3], [2], [1]].
解释：

1

/
2 3
/ \
4 5
样例2
输入： {1,2,3,4}
输出： [[4, 3], [2], [1]].
解释：

1

/
2 3
/
4

/**
 * Definition of TreeNode:
 * public class TreeNode {
 *     public int val;
 *     public TreeNode left, right;
 *     public TreeNode(int val) {
 *         this.val = val;
 *         this.left = this.right = null;
 *     }
 * }
 */


public class Solution {
    /*
     * @param root: the root of binary tree
     * @return: collect and remove all leaves
     */
    int max_depth ;
    Map<Integer , List<Integer>> depth = new HashMap<>() ;
    public List<List<Integer>> findLeaves(TreeNode root) {
        // write your code here
        List<List<Integer>> ans = new ArrayList<>() ;
        max_depth = 0 ;
        dfs(root) ;
        for(int i = 1 ; i <= max_depth ; i++){
            ans.add(depth.get(i)) ;
        }
        return ans ;

    }
    public int dfs(TreeNode cur){
        if(cur == null){
            return 0 ;
        }
        int d = Math.max(dfs(cur.left) , dfs(cur.right)) + 1 ;
        max_depth = Math.max(max_depth , d) ;
        if(!depth.containsKey(d)){
            depth.put(d , new ArrayList<Integer>()) ;
        }
        depth.get(d).add(cur.val) ;
        return d ;
    }
}