Count Univalue Subtrees

, , , ,

Question

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.

Example

1
2
3
4
5
6
7
8
9
   0
  / \
 1   0
    / \
   1   0
  / \
 1   1
 
return 5;

Solution

We need a helper function to tell us if it is unival tree for current root, and when left and right tell us it’s unival tree, we have to check the value of two subtrees.

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
// time: O(n) space:O(n)
public int res = 0;
    public int countUnivalSubtrees(TreeNode root) {
        if (root == null) return 0;
        helper(root);
        return res;
    }
    
    public boolean helper(TreeNode root) {
        if (root == null) return true;
        boolean left = helper(root.left);
        boolean right = helper(root.right);
        if (left && right) {
            if (root.left != null && root.val != root.left.val) return false;
            if (root.right != null && root.val  != root.right.val) return false;
            res++;
            return true;
        }
        return false;
    }