Find Leaves of Binary Tree Python Solution

lcpy gen -n 366   # by problem number
lcpy gen -s find_leaves_of_binary_tree   # by problem name

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.

Input: root = [1]
Output: [[1]]

from leetcode_py import TreeNode


class Solution:
    # Time: O(n) — every node visited once
    # Space: O(h) — recursion depth equals tree height
    def find_leaves(self, root: TreeNode[int] | None) -> list[list[int]]:
        result: list[list[int]] = []

        def height(node: TreeNode[int] | None) -> int:
            if node is None:
                return 0
            h = 1 + max(height(node.left), height(node.right))
            while len(result) < h:
                result.append([])
            result[h - 1].append(node.val)
            return h

        height(root)
        return result