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[];

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

 int treeRoot, flag, blackNodeCount, in[]; // 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 = ; i < m; i++) { // every node has no child
         tree[i].lchild = tree[i].rchild = -;
     }

     flag = ; // at first, it is a red-black tree
     treeRoot = blackNodeCount = ; // 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 + , a1 + i - b1, b1, i - );
     }
     if (i < b2) { // if right subtree exists
         tree[root].rchild = buildTree(a1 + i - b1 + , a2, i + , 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 == ) { // if it is not a red-black tree
         return;
     }

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

         return;
     }

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

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

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

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

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

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

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

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

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

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

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

     system("pause");
     return ;
 }