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:
Given binary tree

          1
         / \
        2   3
       / \     
      4   5

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

Explanation:

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

          1
         / 
        2

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

          1

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

          []

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

Solution

(1) Java

(2) Python

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def findLeaves(self, root):
        """
        :type root: TreeNode
        :rtype: List[List[int]]
        """
        def helper(node):
            if not node:
                return -1
            l = helper(node.left)
            r = helper(node.right)
            level = max(l, r)+1
            if len(rst) == level:
                rst.append([node.val])
            else:
                rst[level].append(node.val)
            return level

        rst =[]
        if not root:
            return rst
        helper(root)
        return rst

(3) Scala