思路
类似与求树的高度(分治)。在每次求得当前节点的高度后,使用hashmap将结果保存下来(key=高度,value=一个list,存的是高度=key的所有节点的value)。由于使用分治法递归求树的高度的过程中,已经将每个节点的高度都求了一遍,所以在每次求得高度后将高度与节点值直接存到hashmap中,就不用对每一个节点进行一遍dfs了,保证时间复杂度为线性。
代码
/**
* 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
*/
public List<List<Integer>> findLeaves(TreeNode root) {
// write your code here
List<List<Integer>> res = new ArrayList<>();
Map<Integer, List<Integer>> hash = new HashMap<>();
int maxHeight = dfs(root, hash);
for(int i = 1; i <= maxHeight; i++) {
res.add(hash.get(i));
}
return res;
}
public int dfs(TreeNode root, Map<Integer, List<Integer>> hash) {
if(root == null) return 0;
// divide (left + right)
int height = Math.max(dfs(root.left, hash), dfs(root.right, hash)) + 1;
hash.putIfAbsent(height, new ArrayList<>());
hash.get(height).add(root.val);
return height;
}
}
复杂度
时间复杂度O(n): 每个节点访问有且仅有一次
空间复杂度O(logn):递归栈