366. Find Leaves of Binary Tree

Problem description:

Given the root of a binary tree, collect a tree’s nodes as if you were doing this:

• Collect all the leaf nodes.
• Remove all the leaf nodes.
• Repeat until the tree is empty.

Example 1:

Input: root = [1,2,3,4,5]
Output: [[4,5,3],[2],[1]]
Explanation:
[[3,5,4],[2],[1]] and [[3,4,5],[2],[1]] are also considered correct answers since per each level it does not matter the order on which elements are returned.

Solution:

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def findLeaves(self, root: TreeNode) -> List[List[int]]:
        def dfs(root):
            if not root:
                return -1
            left = dfs(root.left)
            right = dfs(root.right)
            height = max(left, right) + 1
            dic[height].append(root.val)
            return height
        dic = defaultdict(list)
        dfs(root)
        return dic.values()

time complexity: $O()$
space complexity: $O()$
reference:
related problem: