PAT-A1135. Is It A Red-Black Tree (30)

　　已知先序序列，判断对应的二叉排序树是否为红黑树。序列中负数表示红色结点，正数表示黑色结点。该序列负数取绝对值后再排序得到的是中序序列。根据红黑树的性质判断它是否符合红黑树的要求。考察了根据先序序列和中序序列建树和DFS。

//#include "stdafx.h"
#include <iostream>
#include <algorithm>

using namespace std;

struct treeNode { // tree node struct
    int key, lchild, rchild, flag; // key, left child, right child, color flag
}tree[31];

struct intNode { // int node
    int key, flag; // key, color flag
}pre[31]; // the array of the preorder traversal sequence

int treeRoot, flag, blackNodeCount, in[31]; // the index of tree root node, whether it is a red-black tree, the number of black nodes, the array of the inorder traversal sequence

void init(int m) { // initialize
    int i;
    for (i = 0; i < m; i++) { // every node has no child
        tree[i].lchild = tree[i].rchild = -1;
    }
    
    flag = 1; // at first, it is a red-black tree
    treeRoot = blackNodeCount = 0; // the initial index of tree root node is zero and the initial number of black nodes is zero

    sort(in, in + m); // sort in array to get the inorder traversal sequence
}

int buildTree(int a1, int a2, int b1, int b2) { // build a tree according to the preorder traversal sequence and inorder
    // initialize the current root node of the subtree
    int root = treeRoot++;
    int key = pre[a1].key;
    tree[root].key = key;
    tree[root].flag = pre[a1].flag;

    int i;
    for (i = b1; i <= b2; i++) { // seek the index of root node in the inorder traversal sequence
        if (in[i] == key) {
            break;
        }
    }

    if (i > b1) { // if left subtree exists
        tree[root].lchild = buildTree(a1 + 1, a1 + i - b1, b1, i - 1);
    }
    if (i < b2) { // if right subtree exists
        tree[root].rchild = buildTree(a1 + i - b1 + 1, a2, i + 1, b2);
    }

    return root; // return the current root node of the subtree
}

void dfs(int cur, int count) { // traverse all nodes based on depth first
    if (flag == 0) { // if it is not a red-black tree
        return;
    }

    if (cur == -1) { // if it is a leaf node
        if (count != blackNodeCount) { // judge whether black nodes is legal
            if (blackNodeCount == 0) {
                blackNodeCount = count;
            } else {
                flag = 0;
            }
        }

        return;
    }

    int l = tree[cur].lchild;
    int r = tree[cur].rchild;

    if (tree[cur].flag == 0) { // If a node is red, then both its children are black.
        if (l != -1 && tree[l].flag == 0) {
            flag = 0;
            return;
        } else if (r != -1 && tree[r].flag == 0) {
            flag = 0;
            return;
        }
    } else {
        count++;
    }

    // recursion
    dfs(l, count);
    dfs(r, count);
}

int judgeRBTree() {
    if (tree[0].flag == 0) { // The root is black.
        return 0;
    }

    // traverse
    dfs(0, 0); 
    return flag;
}

int main() {
    int n;
    scanf("%d", &n);

    int m, i;
    while (n--) {
        scanf("%d", &m);
        for (i = 0; i < m; i++) {
            scanf("%d", &pre[i].key);

            // save the information of color
            if (pre[i].key < 0) { 
                pre[i].key = - pre[i].key;
                pre[i].flag = 0;
            } else {
                pre[i].flag = 1;
            }

            in[i] = pre[i].key;
        }

        init(m); // initialization
        buildTree(0, m - 1, 0, m - 1); // build a tree

        if (judgeRBTree() == 1) {
            printf("Yes
");
        } else {
            printf("No
");
        }
    }

    system("pause");
    return 0;
}