Leetcode 250. Count Univalue Subtrees (Medium) (cpp)

Tag: Tree

Difficulty: Medium

/* 
 
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. 
 
For example: 
Given binary tree, 
5 
/ \ 
1   5 
/ \   \ 
5   5   5 
return 4. 
 
*/  
/** 
* Definition for a binary tree node. 
* struct TreeNode { 
*     int val; 
*     TreeNode *left; 
*     TreeNode *right; 
*     TreeNode(int x) : val(x), left(NULL), right(NULL) {} 
* }; 
*/  
class Solution {  
public:  
    int countUnivalSubtrees(TreeNode* root) {  
        int res = 0;  
        cout(root, res);  
        return res;  
    }  
private:  
    int cout(TreeNode* root, int& res) {  
        if (root == NULL) {  
            return 0;  
        }  
        else if (root->left == NULL && root->right == NULL) {  
            res++;  
            return 1;  
        }  
        int l = cout(root->left, res), r = cout(root->right, res);  
        if (l == 1 && r == 1) {  
            if (root->val == root->left->val && root->val == root->right->val) {  
                res += 1;  
                return 1;  
            }  
        }  
        else if (l == -1 || r == -1) {  
            return -1;  
        }  
        else if (l != 0 && root->val == root->left->val) {  
            res += 1;  
            return 1;  
        }  
        else if (r != 0 && root->val == root->right->val) {  
            res += 1;  
            return 1;  
        }  
        return -1;  
    }  
};