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

数据结构中有一类平衡的二叉搜索树，称为红黑树。

它具有以下 5 个属性：

节点是红色或黑色。
根节点是黑色。
所有叶子都是黑色。（叶子是 NULL节点）
每个红色节点的两个子节点都是黑色。
从任一节点到其每个叶子的所有路径都包含相同数目的黑色节点。
例如，下列三张图中，左图中的二叉树是红黑树，其余两图中的二叉树不是红黑树。

现在，对于每个给定的二叉搜索树，请你判断它是否是合法的红黑树。

注意
给定的前序遍历序列可能不合法，即无法构建出合法二叉搜索树。

输入格式
第一行包含整数 K，表示共有 K 组测试数据。

每组测试数据，第一行包含整数 N，表示二叉搜索树的节点数量。

第二行给出了这个二叉搜索树的前序遍历。

注意，虽然所有节点的权值都为正，但是我们使用负号表示红色节点。

各节点权值互不相同。

输入样例与题目中三个图例相对应。

输出格式
对于每组数据，如果是合法红黑树则输出一行 Yes，否则输出一行 No。

数据范围
1≤K≤30,
1≤N≤30
输入样例：
3
9
7 -2 1 5 -4 -11 8 14 -15
9
11 -2 1 -7 5 -4 8 14 -15
8
10 -7 5 -6 8 15 -11 17
输出样例：
Yes
No
No

思路：主体建树与普通的二叉搜索树建树一样，只是多添加一个color标记。
建完过后，在最后一层的叶子再加上2个黑节点
剩下的就是判断了
（1）如果头为红为No
（2）遍历树，如果此节点为red且孩子有一个不为black则为No
（3）层序遍历每个节点，用set保存该节点到每个叶子节点途中black节点的数量，若set的size()>1则为No
（4）上述都筛过成功，为Yes

#include <bits/stdc++.h>

using namespace std;
struct node
{
    int data;
    struct node*l,*r;
    string color;
};
int t,n;
int flag=1;
int ll,rr;
set<int> s;
struct node*create(string color,int shu,struct node*root)
{
    if(root == NULL)
    {
        root = new node;
        root->data =  shu;
        root->color = color;
        root->l=NULL;
        root->r=NULL;
        return root;
    }
    else if(shu<root->data)
    {
        root->l = create(color,shu,root->l);
    }
    else
      root->r = create(color,shu,root->r);
      
      return root;
};
void add(struct node*&root)
{
    if(root)
    {
        add(root->l);
        add(root->r);
    }
    else
    {
        root = new node();
        root->data = 0;
        root->color = "black";
        root->l=NULL;
        root->r=NULL;
        return ;
    }
}
void  judge(struct node*root)
{
    if(root)
    {
        judge(root->l);
        judge(root->r); 
        if(root->color=="red"&&(root->l->color!="black"||root->r->color!="black"))
            flag=0;
    }
}
int shendu(struct node*root,int num)
{
   if(root)
   {
       if(root->color == "black")
       {
           num++;
       }
       shendu(root->l,num);
       shendu(root->r,num);
   }
   else 
   s.insert(num);
}
void cengxv(struct node*root)
{
  queue<node*> q;
  q.push(root);
  while(!q.empty())
  {
      node* now = q.front();
      s.clear();
      shendu(now,0);
      if(s.size()>1)
      flag=0;
      q.pop();
      if(now->l)
      q.push(now->l);
      if(now->r)
      q.push(now->r);
  }
}
int main()
{
    cin >> t;
    while(t--)
    {
        flag=1;
        cin >> n;
        getchar();
        struct node*root = NULL;
        for(int i=1;i<=n;i++)
        {
             string s;
             cin >>s;
           
             int x = stoi(s);  
             if(x<0)
             root = create("red",-x,root);
             else
             root = create("black",x,root);
        }
        add(root);
        if(root->color!="black")
        {
              cout <<"No"<<endl;
           continue;
        }
       judge(root);
      if(!flag)
      {
           cout <<"No"<<endl;
           continue;
      }
      cengxv(root);
      if(!flag)
      {
           cout <<"No"<<endl;
           continue;
      }
      cout <<"Yes"<<endl;
    }
    return 0;
}