Tree - 250. Count Univalue Subtrees

6

250. Count Univalue Subtrees

Given a binary tree, count the number of uni-value subtrees.
A Uni-value subtree means all nodes of the subtree have the same value.

Example

Example1

Input:  root = {5,1,5,5,5,#,5}
Output: 4
Explanation:
              5
             / \
            1   5
           / \   \
          5   5   5 

Example2

Input:  root = {1,3,2,4,5,#,6}
Output: 3
Explanation:
              1
             / \
            3   2
           / \   \
          4   5   6

思路:

题目意思是找出一棵树中以树中某节点为根节点的子树中所有元素都相同,很明显,叶子结点肯定是一个uni的子树,回到问题本身,判断以一个节点为根节点的子树是否为uni-value,只需要满足左子树是uni-value,右子树是uni-value,并且判断当前节点和左右孩子节点值,就能求解。所以,递归后序遍历整个树,就能得出结果,每一次递归返回,返回的是以当前节点为根节点是否是uni-value,其中corner case,就是节点为nil,默认为uni-value,返回给上一层就可以。

代码:

go:

/**
 * 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 {
    var res = 0
    if root != nil {
        _ = isUnivalSubtrees(root, &res)
    }
    return res
}

func isUnivalSubtrees(node *TreeNode, res *int) bool {
    if node == nil {
        return true
    }
    
    left := isUnivalSubtrees(node.Left, res)
    right := isUnivalSubtrees(node.Right, res)
    
    if left && right && 
        (node.Left == nil || node.Val == node.Left.Val) &&
        (node.Right == nil || node.Val == node.Right.Val) {
        *res++
        return true
    }
    
    return false
}