Counting Unival Subtrees

   0
  / \
 1   0
    / \
   1   0
  / \
 1   1

1
1
1
   1  
  / \
 1   1
 0

/**
* Definition for a binary tree node.
* type TreeNode struct {
*     Val int
*     Left *TreeNode
*     Right *TreeNode
* }
*/

/**
* @param root: the given tree
* @return: the number of uni-value subtrees.
*/
func countUnivalSubtrees (root *TreeNode) int {
    count, _ := _countUnivalSubtrees(root)
    return count
}

func _countUnivalSubtrees (root *TreeNode) (count int, isUnival bool) {
    if root == nil {
    return 0, true
    }
    if root.Left == nil && root.Right == nil {
        return 1, true
    }
    cl, isUnivalL := _countUnivalSubtrees(root.Left)
    cr, isUnivalR := _countUnivalSubtrees(root.Right)
    if isUnivalL && isUnivalR {
        if root.Left == nil && root.Val == root.Right.Val {
            return cl + 1 + cr, true
        } else if root.Right == nil && root.Val == root.Left.Val {
            return cl + 1 + cr, true
        } else if root.Left != nil && root.Right != nil && root.Val == root.Left.Val && root.Val == root.Right.Val {
            return cl + 1 + cr, true
        }
    }
    return cl + cr, false
}