Tuesday, September 8, 2015

Leetcode: 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.
For example:
Given binary tree,
              5
             / \
            1   5
           / \   \
          5   5   5
return 4.
Understand the problem:
The problem can be solved by divide-and-conquer. The only trick is to use different return values to mark different cases. 
Integer.MIN_VALUE -- Mark the subtree is not univaled. 
Integer.MAX_VALUE -- Mark if the root is null. 

Code (Java):
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    private int count = 0;
    public int countUnivalSubtrees(TreeNode root) {
        if (root == null) {
            return 0;
        }
        
        countUnivalSubtreesHelper(root);
        
        return count;
    }
    
    private int countUnivalSubtreesHelper(TreeNode root) {
        if (root == null) {
            return Integer.MAX_VALUE;
        }
        
        // Divide
        int left = countUnivalSubtreesHelper(root.left);
        int right = countUnivalSubtreesHelper(root.right);
        
        // left and right are all empty
        if (left == Integer.MAX_VALUE && right == Integer.MAX_VALUE) {
            count++;
            return root.val;
        } else if (left == Integer.MAX_VALUE || right == Integer.MAX_VALUE) {
            if (root.val == left || root.val == right) {
                count++;
                return root.val;
            } else {
                return Integer.MIN_VALUE;
            }
        } else {
            if (root.val == left && root.val == right) {
                count++;
                return root.val;
            } else {
                return Integer.MIN_VALUE;
            }
        }
    }
}

No comments:

Post a Comment