[LintCode/LeetCode] Count Univalue Subtrees

Problem

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

Given root = {5,1,5,5,5,#,5}, return 4.

          5
         / \
        1   5
       / \   \
      5   5   5
     

Solution

public class Solution {
    /**
     * @param root: the given tree
     * @return: the number of uni-value subtrees.
     */
    private int count = 0;
    public int countUnivalSubtrees(TreeNode root) {
        // write your code here
        helper(root, 0);
        return count;
    }

    private boolean helper(TreeNode root, int value) {
        if (root == null) return true;
        // we need the root.val as the reference in the recursion
        boolean left = helper(root.left, root.val);
        boolean right = helper(root.right, root.val);
        if (left && right) {
            count++;
            return root.val == value;
        }
        
        return false;
    }
}