https://leetcode.com/problems/count-univalue-subtrees

countUnivalSubtrees.py

def isUnival(root, val):
    if root is None:
        return True
    if root.val != val:
        return False

    left = isUnival(root.left, val)
    right = isUnival(root.right, val)

    if not left or not right:
        return False

    return True


class Solution:
    def countUnivalSubtrees(self, root: Optional[TreeNode]) -> int:
        count = 0
        
        def helper(root):
            nonlocal count
            if root is None:
                return 0
            
            print(root.val, isUnival(root, root.val))

            if isUnival(root, root.val):
                print(True)
                count += 1

            helper(root.left) 
            helper(root.right)

            return count
        
        helper(root)

        return count
        

countUnivalSubtrees_optimized.py

class Solution:
    def countUnivalSubtrees(self, root: Optional[TreeNode]) -> int:
        
        def helper(root):
            if root is None:
                return [0, True]
            
            leftCount, isLeftUnival = helper(root.left)
            rightCount, isRightUnival = helper(root.right)


            isUniVal = True
            if not isLeftUnival or not isRightUnival:
                isUniVal = False
            if root.left and root.left.val != root.val:
                isUniVal = False
            if root.right and root.right.val != root.val:
                isUniVal = False

            if isUniVal:
                return [leftCount + rightCount + 1, isUniVal]
            return [leftCount + rightCount, isUniVal]
            

        [count, isUnival] = helper(root)

        return count