Problem statement

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

Solution

Traverse tree with dfs(node) function, which will return pair of elements: first one is True if subtree is univalue and second one is value of this subtree. We also use None value for nodes outside tree. How we can check now if subtree is univalue? We check answer for left and for right. If both of them are True ans also if left and right value both in [None, node.val], then we are happy: we add 1 to final answer.

Complexity

Time complexity is $O(n)$, where $n$ is number of nodes, space complexity is $O(h)$, where $h$ is heigth of the tree.

Code

class Solution:
    def countUnivalSubtrees(self, root):
        self.ans = 0
        
        def dfs(node):
            if not node: return (True, None)
            L, R = dfs(node.left), dfs(node.right)
            cands = [None, node.val]
            
            found = L[0] and R[0] and R[1] in cands and L[1] in cands
            if found: self.ans += 1
            return (found, node.val)
        
        dfs(root)
        return self.ans