[LintCode] Problem 650 - Find Leaves of Binary Tree

Posted on | In ,

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

No.1

Input: {1,2,3,4,5}

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

Explanation:

1
2
3
4
5
    1
   / \
  2   3
 / \     
4   5

No.2

Input: {1,2,3,4}

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

Explanation:

1
2
3
4
5
    1
   / \
  2   3
 /
4

Code

1
2
3
4
5
6
7
8
public class TreeNode {
    public int val;
    public TreeNode left, right;
    public TreeNode(int val) {
        this.val = val;
        this.left = this.right = null;
    }
}
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
public List<List<Integer>> findLeaves(TreeNode root) {
    Map<Integer, List<Integer>> map = new HashMap<>();
    List<List<Integer>> result = new ArrayList<>();

    findLeavesHelper(root, map);

    for (List<Integer> values : map.values())
        result.add(values);

    return result;
}

private int findLeavesHelper(TreeNode root, Map<Integer, List<Integer>> map) {
    if (root == null)
        return 0;

    int left = findLeavesHelper(root.left, map);
    int right = findLeavesHelper(root.right, map);

    int level = Math.max(left, right) + 1;
    map.putIfAbsent(level, new ArrayList<>());
    map.get(level).add(root.val);

    return level;
}