题目来源

• leetcode：250-vip. 统计同值子树

题目描述

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};

class Solution{
    int countUnivalSubtrees(TreeNode* root){
        
    }
};

题目解析

思路

• 首先，叶子节点都是同值子树
• 然后，对于非叶子节点：
  • 如果仅仅存在左子树，那么必须左子树是同值子树&&左子树的节点值必须等于当前节点的值
  • 如果仅仅存在右子树，那么必须右子树是同值子树&&右子树的节点值必须等于当前节点的值
  • 如果存在左右子树，那么必须左右子树都是同值子树&&左子树的节点值必须等于当前节点的值&&右子树的节点值必须等于当前节点的值

实现

class Solution{
    struct Info{
        bool isUnival;
        int cnt;
        int val;

        Info(bool isUnival, int cnt, int val) : isUnival(isUnival), cnt(cnt), val(val){

        }
    };

    Info *process(TreeNode* root){
        if(root == nullptr){
            return nullptr;
        }

        auto left = process(root->left);
        auto right = process(root->right);
        int cnt = 0;
        bool isUnival = false;
        if(left == nullptr && right == nullptr){
            isUnival = true;
            cnt = 1;
        }else if(left == nullptr){
            isUnival = right->isUnival && right->val == root->val;
            cnt = (isUnival? 1 : 0) + right->cnt;
        }else if(right == nullptr){
            isUnival = left->isUnival && left->val == root->val;
            cnt = (isUnival? 1 : 0) + left->cnt;
        }else{
            isUnival = right->isUnival && right->val == root->val && left->isUnival && left->val == root->val;
            cnt = (isUnival? 1 : 0) + left->cnt + right->cnt;
        }
        return new Info(isUnival, cnt, root->val);
    }
public:
    int countUnivalSubtrees(TreeNode* root){
        if(root == nullptr){
            return 0;
        }
        return process(root)->cnt;
    }
};

测试

int main() {
    TreeNode* root = new TreeNode(5);
    root->left = new TreeNode(1);
    root->right = new TreeNode(5);
    root->left->left = new TreeNode(5);
    root->left->right = new TreeNode(5);
    root->right->right = new TreeNode(5);

    Solution a;
    std::cout << a.countUnivalSubtrees(root) << "\n";
    return 1;
}