366. Find Leaves of Binary Tree (Medium) - TengnanYao/daily_leetcode GitHub Wiki

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution(object):
    def findLeaves(self, root):
        """
        :type root: TreeNode
        :rtype: List[List[int]]
        """
        result = []
        while root.left or root.right:
            temp = []
            arr = [root]
            while arr:
                nodes = []
                for node in arr:
                    if node.left:
                        if not node.left.left and not node.left.right:
                            temp.append(node.left.val)
                            node.left = None
                        else:
                            nodes.append(node.left)
                    if node.right:
                        if not node.right.left and not node.right.right:
                            temp.append(node.right.val)
                            node.right = None
                        else:
                            nodes.append(node.right)
                arr = nodes
            result.append(temp)
        result.append([root.val])
        return result