PAT 1135 Is It A Red-Black Tree

There is a kind of balanced binary search tree named red-black tree in the data structure. It has the following 5 properties:

• (1) Every node is either red or black.
• (2) The root is black.
• (3) Every leaf (NULL) is black.
• (4) If a node is red, then both its children are black.
• (5) For each node, all simple paths from the node to descendant leaves contain the same number of black nodes.

For example, the tree in Figure 1 is a red-black tree, while the ones in Figure 2 and 3 are not.

Figure 1	Figure 2	Figure 3

For each given binary search tree, you are supposed to tell if it is a legal red-black tree.

Input Specification:

Each input file contains several test cases. The first line gives a positive integer K (≤30) which is the total number of cases. For each case, the first line gives a positive integer N (≤30), the total number of nodes in the binary tree. The second line gives the preorder traversal sequence of the tree. While all the keys in a tree are positive integers, we use negative signs to represent red nodes. All the numbers in a line are separated by a space. The sample input cases correspond to the trees shown in Figure 1, 2 and 3.

Output Specification:

For each test case, print in a line "Yes" if the given tree is a red-black tree, or "No" if not.

Sample Input:

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

Sample Output:

Yes
No
No

主要关键点在于判断到达每个叶子节点有数量相同的黑节点。换一种想法，就是最大值和最小值相同。判断每一个节点左边路径上最大数量黑节点和有变路径最小数量黑节点是否相同。

#include <iostream>
#include<vector>
#include<set>
#include<cmath>
using namespace std;
struct TreeNode{
    int val{};
    TreeNode* left{};
    TreeNode* right{};
    bool isBlack=true;
};
vector<int> v1;
TreeNode* constructTree(int beg,int end){
    if(beg>end){
        return nullptr;
    }else if(beg==end){
        int val;
        auto* node=new TreeNode();
        if(v1[beg]<0){
            val=-v1[beg];
            node->isBlack=false;
        }else{
            val=v1[beg];
        }
        node->val=val;
        return node;
    }
    int val=abs(v1[beg]);
    int index=beg+1;
    for(int i=beg+1;i<=end;i++){
        if(abs(v1[i])>val){
            index=i;
            break;
        }
    }
    auto* node=new TreeNode();
    node->val=val;
    if(v1[beg]<0){
        node->isBlack=false;
    }
    node->left=constructTree(beg+1,index-1);
    node->right=constructTree(index,end);
    return node;
}
int countBlk1(TreeNode* node){
    if(node==nullptr){
        return 1;
    }
    int left=countBlk1(node->left);
    int right=countBlk1(node->right);
    if(node->isBlack){
        return max(left,right)+1;
    }
    return max(left,right);
}
int countBlk2(TreeNode* node){
    if(node==nullptr){
        return 1;
    }
    int left=countBlk2(node->left);
    int right=countBlk2(node->right);
    if(node->isBlack){
        return min(left,right)+1;
    }
    return min(left,right);
}
bool fun1(TreeNode* node){
    if(node==nullptr){
        return true;
    }
    if(!node->isBlack){
        if(node->left!=nullptr && !node->left->isBlack){
            return false;
        }
        if(node->right!=nullptr && !node->right->isBlack){
            return false;
        }
    }
    if(countBlk1(node->left)!=countBlk2(node->right)){
        return false;
    }
    return fun1(node->left)&&fun1(node->right);
}
bool judge(TreeNode* root){
    if(!root->isBlack){
        return false;
    }
    return fun1(root);
}
int main()
{
    int K;
    cin>>K;
    vector<bool> vec;
    for(int i=0;i<K;i++){
        int num;
        cin>>num;
        v1.clear();
        for(int j=0;j<num;j++){
            int val;
            cin>>val;
            v1.push_back(val);
        }
        TreeNode* root=constructTree(0,num-1);
        vec.push_back(judge(root));
    }
    for(auto && i : vec){
        if(i){
            cout<<"Yes"<<endl;
        }else{
            cout<<"No"<<endl;
        }
    }
    return 0;
}