pat1135 Is It A Red-Black Tree

题意：给一个树判断是不是红黑树。

思路：根据5条性质判断就好。第5条根据每个节点dfs就好。

代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>

using namespace std;

const int MAX_N = 50;
int T, N, root;
int a[MAX_N], lc[MAX_N], rc[MAX_N];

int build(int lo, int hi) {
    if (lo >= hi) return -1;
    if (hi-lo==1) return lo;
    int p = hi;
    for (int i = lo+1; i < hi; i++) {
        if (abs(a[i]) > abs(a[lo])) {
            p = i; break;
        }
    }
    lc[lo] = build(lo+1, p);
    rc[lo] = build(p, hi);
    return lo;
}

bool flag;
int cnt;
void check(int p, int ct) {
    if (p == -1) return ;
    if (!flag) return ;
    if (a[p] < 0) {
        if (lc[p] != -1 && a[lc[p]] < 0) {
            flag = false; return ;
        }
        if (rc[p] != -1 && a[rc[p]] < 0) {
            flag = false; return ;
        }
    }
    if (a[p] > 0) ct++;
    if (lc[p] == -1 || rc[p] == -1) {
        if (cnt == -1) cnt = ct + 1;
        else if (cnt != ct + 1) {
            flag = false; return ;
        }
    }
    check(lc[p], ct); check(rc[p], ct);
}
void dfs(int p) {
    if (!flag) return ;
    if (p == -1) return ;
    cnt = -1;
    check(p, 0);
    dfs(lc[p]);
    dfs(rc[p]);
}

int main() {
    // freopen("in.txt", "r", stdin);
    // freopen("out.txt", "w", stdout);
    scanf("%d", &T);
    while (T--) {
        scanf("%d", &N); flag = true;
        memset(lc, -1, sizeof(lc));
        memset(rc, -1, sizeof(rc));
        for (int i = 0; i < N; i++) scanf("%d", &a[i]);
        root = build(0, N);
        if (a[root] < 0) {
            printf("No\n"); continue;
        }
        dfs(root);
        if (flag) {
            printf("Yes\n");
        } else {
            printf("No\n");
        }
    }
    return 0;
}