[leetcode] Count Univalue Subtrees


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.

/**
 * 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) {
        if(root == nullptr) return 0;
        int res = 0;
        bool isUni = true;
        DFS(root, res, isUni);
        return res;
    }
    //return uni value, res stores the num of univalue substrees.
    int DFS(TreeNode * root, int & res, bool & isUni){
        //leaf node
        if(root->left == nullptr && root->right == nullptr){
            isUni = true;
            res++;
            return root->val;
        }
        bool isLeftUni = true, isRightUni = true;
        int leftval, rightval;
        if(root->left){
            leftval = DFS(root->left, res, isLeftUni);
        }
        if(root->right){
            rightval = DFS(root->right, res, isRightUni);
        }
        if(isLeftUni && isRightUni){
            leftval = root->left? leftval: root->val;
            rightval = root->right? rightval: root->val;
            isUni = (leftval == rightval) && (root->val == leftval);
            if(isUni == true){
                res++;
            }
            return root->val;
        }
        else{
            //left or right is not uni
            isUni = false;
            return root->val;//useless
        }
    }
};

 

Leave a comment

Your email address will not be published. Required fields are marked *

This site uses Akismet to reduce spam. Learn how your comment data is processed.