BinarySearch 0294 Partition Tree

Problem statement

https://binarysearch.com/problems/Partition-Tree/

Solution

Just traverse tree with dfs/bfs.

Complexity

It is O(n) for time and O(h) for space.

Code

class Solution:
    def solve(self, root):
        self.ans = [0, 0]
        def dfs(node):
            if not node.left and not node.right:
                self.ans[0] += 1
            else:
                self.ans[1] += 1
            if node.left: dfs(node.left)
            if node.right: dfs(node.right)
        
        if not root: return self.ans
        dfs(root)
        return self.ans