250. Count Univalue Subtrees - Medium

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 :

Input:  root = [5,1,5,5,5,null,5]

              5
             / \
            1   5
           / \   \
          5   5   5

Output: 4

 

两次recursion，如果root是univalue的，返回1+ left + right，如果不是，返回left + right（有大量重复check，慢）

time: O(n^2), space: O(height)

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {    
    public int countUnivalSubtrees(TreeNode root) {
        if(root == null) {
            return 0;
        }
        if(unival(root)) {
            return 1 + countUnivalSubtrees(root.left) + countUnivalSubtrees(root.right);
        }
        return countUnivalSubtrees(root.left) + countUnivalSubtrees(root.right);
    }
    
    public boolean unival(TreeNode root) {
        if(root == null) {
            return true;
        }
        if(root.left == null && root.right == null) {
            return true;
        }
        if(root.left != null && root.val != root.left.val) {
            return false;
        }
        if(root.right != null && root.val != root.right.val) {
            return false;
        }
        
        return unival(root.left) && unival(root.right);
    }
}

 

optimized: one pass

time: O(n), space: O(height)

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    int cnt = 0;
    
    public int countUnivalSubtrees(TreeNode root) {
        unival(root);
        return cnt;
    }
    
    public boolean unival(TreeNode root) {
        if(root == null) {
            return true;
        }
        boolean left = unival(root.left);
        boolean right = unival(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;
            }
            cnt++;
            return true;
        }
        return false;
    }
}