【PAT】1135 Is It A Red-Black Tree (30 分)

#include <iostream>
#include <vector>
#include <math.h>
using namespace std;

vector<int> arr;
struct node
{
    int val;
    node *left,*right;
};

node* buildT(node* root,int v){
    if(!root){
        root=new node;
        root->val=v;
        root->left=root->right=NULL;
        return root;
    }
    if(abs(root->val)>abs(v)){
        root->left=buildT(root->left,v);
    }else{
        root->right=buildT(root->right,v);
    }
    return root;
}

bool judge1(node* root){//1：判断子结点颜色是否符合
    if(root==NULL) return true;//叶子或空树——前面已经判断过,退出
    
    if(root->val<0){
        if(root->left!=NULL && root->left->val<0) return false;//不满足,提前退出
        if(root->right!=NULL && root->right->val<0) return false;//不满足,提前退出
    }
//递归判断子树
    return judge1(root->left) && judge1(root->right);
}

int getNum(node* root){//找黑结点最多的路径
    if(!root) return 0;
    int l=getNum(root->left);
    int r=getNum(root->right);
    return (root->val>0)? max(l,r)+1:max(l,r);
}

bool judge2(node* root){//2：判断子结点数量是否符合
    if(!root) return true;
    int l=getNum(root->left);
    int r=getNum(root->right);
    if(l!=r) return false;//递归式放前面,从后往前递归
    return judge2(root->left) && judge2(root->right);//递归式放后面,从前往后递归
}

int main(){
    int k,n;
    scanf("%d",&k);

    while(k--){
        scanf("%d",&n);
        arr.resize(n);
        node* root=NULL;
        for(int i=0;i<n;i++){
            scanf("%d",&arr[i]);
            root=buildT(root,arr[i]);
        }
        if(arr[0]<0 || !judge1(root) || !judge2(root)){
            printf("No\n");
        }else{
            printf("Yes\n");
        }

    }
    return 0;
}