Count Univalue Subtrees

Problem Id: 250 Difficulty: Medium Tag: Tree


Intuition

Using DFS and with unique value of the dfs's return value.

Solution


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

class Solution:
    EMPTY = 'empty'
    NOTUNIQUE = 'not-unique'

    def countUnivalSubtrees(self, root: TreeNode) -> int:
        self.count = 0
        self._helper(root)
        return self.count

    def _helper(self, node):
        if node is None:
            return self.EMPTY
        l = self._helper(node.left)
        r = self._helper(node.right)
        if self.NOTUNIQUE in [l, r]:
            return self.NOTUNIQUE
        if l != self.EMPTY and l != node.val:
            return self.NOTUNIQUE
        if r != self.EMPTY and r != node.val:
            return self.NOTUNIQUE
        self.count += 1
        return node.val