Find Leaves of Binary Tree – Python 3 (Week 12 – 18)

"""
Definition of TreeNode:
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left, self.right = None, None
"""


class Solution:
    """
    @param: root: the root of binary tree
    @return: collect and remove all leaves
    """
    def __init__(self):
        self.ans = []
        
    def findLeaves(self, root):
        # write your code here
        self.dfs(root)
        return self.ans
        
        
    
    def dfs(self, root):
        
        if root == None:
            return -1
        
        left_height = self.dfs(root.left)
        right_height = self.dfs(root.right)
        height = 1 + max(left_height, right_height)
        if height >= len(self.ans):
            self.ans.append([])
        self.ans[height].append(root.val)
        return height