Daily Coding Problem - 8

day-8.py

"""
This problem was asked by Google.

A unival tree (which stands for "universal value") is a tree where all nodes under it have the same value.

Given the root to a binary tree, count the number of unival subtrees.

For example, the following tree has 5 unival subtrees:

   0
  / \
 1   0
    / \
   1   0
  / \
 1   1
"""
class Node:
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def countUnivals(root):
    def helper(root):
        if root.left is None and root.right is None:
            return 1, True
        left_n_univals, left_is_unival = helper(root.left)
        right_n_univals, right_is_unival = helper(root.right)
        if root.left.val == root.right.val:
            if left_is_unival and right_is_unival:
                return 1 + left_n_univals + right_n_univals, True
        return left_n_univals + right_n_univals, False
    count, flag = helper(root)
    return count

binaryTree = Node(0, Node(1), Node(0, Node(1, Node(1), Node(1)), Node(0)))

print(countUnivals(binaryTree))