Count Univalue Subtrees

Problem: Count Univalue Subtrees

When we check a tree is univalue or not, we need to check its subtrees first. Thus, we can do dfs to the tree and once we found a subtree, we add it to a global count. Besides, we also tell parent of root of subtrees so that parent can decide whether itself is a root of a subtree.

Code in Python:

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def countUnivalSubtrees(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        self.ans = 0
        if not root: return self.ans
        def dfs(node):
            if not node.left and not node.right:
                self.ans += 1
                return True
            left = right = None
            if node.left: left = dfs(node.left)
            if node.right: right = dfs(node.right)
            if left == False or right == False: return False
            if not node.left:
                if node.right.val == node.val:
                    self.ans += 1
                    return True
                else:
                    return False
            if not node.right:
                if node.left.val == node.val:
                    self.ans += 1
                    return True
                else:
                    return False
            if node.left.val == node.right.val and node.left.val == node.val:
                self.ans += 1
                return True
            return False
        dfs(root)
        return self.ans