1135 Is It A Red-Black Tree

#include <cstdio>
#include <algorithm>
using namespace std;
#define BLACK 1
#define RED 0

struct Node{
    int val;
    int color;
    Node *lchild,*rchild;
    Node(int v){
        val=v>0?v:-v;
        color=v>0?BLACK:RED;
        lchild=rchild=nullptr;
    }
};

void insert(Node* &root,int val)
{
    if(root==nullptr){
        root=new Node(val);
        return;
    }
    if(abs(val)<root->val) insert(root->lchild,val);
    else insert(root->rchild,val);
}

//深度遍历，计算黑色结点的个数以及判断是否会出现连续两个红色结点
int totalBlackCnt=0;
bool flag=true;
void dfs(Node* root,int cnt)
{
    if(root==nullptr){
        if(cnt!=totalBlackCnt) flag=false;
        return;
    }
    if(root->color==BLACK) cnt++;
    else{
        if(root->lchild && root->lchild->color==RED) flag=false;
        if(root->rchild && root->rchild->color==RED) flag=false;
    }
    //printf("val:%d color:%d cnt:%d\n",root->val,root->color,cnt);

    dfs(root->lchild,cnt);
    dfs(root->rchild,cnt);
}

int main()
{int k,n,val;
    scanf("%d",&k);
    while(k--){
        scanf("%d",&n);
        Node* root=nullptr;
        for(int i=0;i<n;i++){
            scanf("%d",&val);
            insert(root,val);
        }
        if(root->color==RED){
            printf("No\n");
            continue;
        }
        totalBlackCnt=0;//记录从根结点到任意一个叶结点的简单路径上黑色结点的个数
        //此处计算最左端的路径
        Node* p=root;
        while(p){
            if(p->color==BLACK) totalBlackCnt++;
            p=p->lchild;
        }
        flag=true;//初始化
        dfs(root,0);
        printf("%s\n",flag?"Yes":"No");
    }
    return 0;
}

void dfs(Node* root,int cnt)
{
    if(root==nullptr) return;
    if(root->color==BLACK) cnt++;
    else{
        if(root->lchild && root->lchild->color==RED) flag=false;
        if(root->rchild && root->rchild->color==RED) flag=false;
    }
    if(root->lchild==nullptr && root->rchild==nullptr){
        if(cnt!=totalBlackCnt) flag=false;
    }
    //printf("val:%d color:%d cnt:%d\n",root->val,root->color,cnt);
    dfs(root->lchild,cnt);
    dfs(root->rchild,cnt);
}