Problem statement

https://leetcode.com/problems/find-leaves-of-binary-tree/

Solution

What we actually need is to evaluate level of each node. We can use dfs for it. If we have None node, level is 0. Else, we evaluate maximum of levels of children and add $1$. I put nodes in defaultdict(list) and then create list.

Complexity

Time and space complexity is $O(n)$

Code

class Solution:
    def findLeaves(self, root):
        def dfs(root):
            if not root: return 0

            left_cand = dfs(root.left)
            right_cand = dfs(root.right)
            lev = max(left_cand, right_cand) + 1
            levels[lev].append(root.val)
            return lev

        levels = defaultdict(list)
        dfs(root)
        return [levels[i+1] for i in range(len(levels))]