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

 

#include<cstdio>
#include<vector>
#define maxn 33
#define inf 300
using namespace std;

struct simplenode {
    int d;
    bool red;
    simplenode(int _d, bool _red):d(_d), red(_red) {}
};

struct node {
    int d;
    bool red;
    node *left, *right;
};

node* buildtree(vector<simplenode> in, int sta, int en) {
    if (sta >= en) return NULL;
    int i;
    for (i = sta + 1; i < en; i++) {
        if (in[i].d >= in[sta].d) break;
    }
    node *root = new node;
    root->d = in[sta].d;
    root->red = in[sta].red;
    root->left = buildtree(in, sta + 1, i);
    root->right = buildtree(in, i, en);
    return root;
}

// return the num of the black nodes the tree base on "root" have,
// while checking it is valid
int check_black_num (node *root) {
    if (root == NULL) return 1;

    int left = check_black_num(root->left);
    if (left >= inf) return inf;
    int right = check_black_num(root->right);
    if (right >= inf) return inf;
    if (left != right) return inf;

    if (root->red == false) return left + 1;
    else return left;
}

bool check_children_red(node *root) {
    if (root == NULL) return true;
    if (root->red == true) {
        if (root->left != NULL && root->left->red == true) return false;
        if (root->right != NULL && root->right->red == true) return false;
    }
    if (root->left != NULL && check_children_red(root->left) == false) return false;
    if (root->right != NULL && check_children_red(root->right) == false) return false;
    return true;
}

bool check_red_black(node *root) {
    if (root->red == true) return false;
    if (check_children_red(root) == false) return false;
    int left = check_black_num(root->left);
    if (left >= inf) return false;
    int right = check_black_num(root->right);
    if (right >= inf) return false;
    if (left != right) return false;
    return true;
}

int main() {
    int K, N, a;
    scanf("%d", &K);
    while(K--) {
        scanf("%d", &N);
        vector<simplenode> input;
        for (int i = 0; i < N; i++) {
            scanf("%d", &a);
            if (a < 0) {
                input.push_back(simplenode(-a, true));
            }
            else input.push_back(simplenode(a, false));
        }
        node *root = buildtree(input, 0, N);
        if (check_red_black(root)) printf("Yes\n");
        else printf("No\n");

    }
    return 0;
}