Problem statement

https://binarysearch.com/problems/Partition-Zero-One-Trees/

Solution

First, run dfs to calculate total number of ones and zeroes. Then run one more dfs to check for given subtree if it has all ones and no zeroes or all zeroes and no ones.

Complexity

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

Code

class Solution:
    def solve(self, root):
        def dfs1(node):
            if not node: return
            d[node.val] += 1
            dfs1(node.left)
            dfs1(node.right)
        
        d = Counter()
        dfs1(root)

        def dfs(node):
            if not node: return (0, 0)
            z0, o0 = dfs(node.left)
            z1, o1 = dfs(node.right)
            z2, o2 = z0 + z1 + int(node.val == 0), o0 + o1 + int(node.val == 1)
            if (o2 == d[1] and z2 == 0) or (o2 == 0 and z2 == d[0]):
                self.ans += 1
            return z2, o2
        
        self.ans = 0
        dfs(root)
        return self.ans