Count Univalue Subtrees Python Solution

lcpy gen -n 250   # by problem number
lcpy gen -s count_univalue_subtrees   # by problem name

Input: root = [5,1,5,5,5,null,5]
Output: 4

Input: root = []
Output: 0

Input: root = [5,5,5,5,5,null,5]
Output: 6

from leetcode_py import TreeNode


class Solution:
    # Time: O(n)
    # Space: O(n)
    def count_unival_subtrees(self, root: TreeNode[int] | None) -> int:
        count = 0

        def is_unival(node: TreeNode[int] | None) -> bool:
            nonlocal count
            if node is None:
                return True
            left_unival = is_unival(node.left)
            right_unival = is_unival(node.right)
            if not left_unival or not right_unival:
                return False
            if node.left is not None and node.left.val != node.val:
                return False
            if node.right is not None and node.right.val != node.val:
                return False
            count += 1
            return True

        is_unival(root)
        return count