【PTA数据结构 | C语言版】是否红黑树

本专栏持续输出数据结构题目集，欢迎订阅。

题目

给定一棵染了红/黑颜色的二叉查找树，请判定其是否满足红黑树定义。

输入格式：
输入首先给出整数 k（≤30），是测试数据的组数。每一组测试数据先在一行中给出正整数 n（≤30），为树中结点个数；第二行给出该二叉查找树树的先序遍历序列。树中所有键值都是正整数，但我们用负号表示该对应结点是红色的，其它没有负号的结点是黑色的。同行数字以空格分隔。

输出格式：
对每组数据，如果对应的二叉查找树是红黑树，就在一行中输出 Yes，否则输出 No。

输入样例：
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

输出样例：
Yes
No
No

代码

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <limits.h>

// 树节点结构
typedef struct Node {
    int key;           // 键值（绝对值）
    bool is_red;       // 是否为红色节点
    struct Node *left;  // 左子节点
    struct Node *right; // 右子节点
} Node;

// 创建新节点
Node* createNode(int key, bool is_red) {
    Node* node = (Node*)malloc(sizeof(Node));
    node->key = abs(key);  // 存储键值的绝对值
    node->is_red = is_red;
    node->left = NULL;
    node->right = NULL;
    return node;
}

// 插入节点到二叉查找树
Node* insert(Node* root, int key, bool is_red) {
    if (root == NULL) {
        return createNode(key, is_red);
    }
    
    int abs_key = abs(key);
    if (abs_key < root->key) {
        root->left = insert(root->left, key, is_red);
    } else {
        root->right = insert(root->right, key, is_red);
    }
    return root;
}

// 验证二叉查找树性质
bool isBST(Node* node, int min_val, int max_val) {
    if (node == NULL) return true;
    
    // 检查当前节点是否在[min_val, max_val]范围内
    if (node->key < min_val || node->key > max_val) {
        return false;
    }
    
    // 递归检查左右子树
    return isBST(node->left, min_val, node->key - 1) && 
           isBST(node->right, node->key + 1, max_val);
}

// 计算从节点到叶子的黑色节点数
int countBlackHeight(Node* node) {
    if (node == NULL) return 1;  // 空节点算一个黑色节点
    
    int left_black = countBlackHeight(node->left);
    int right_black = countBlackHeight(node->right);
    
    // 如果左右子树黑色高度不一致，返回-1表示无效
    if (left_black == -1 || right_black == -1 || left_black != right_black) {
        return -1;
    }
    
    // 当前节点如果是黑色，高度加1
    return (node->is_red ? 0 : 1) + left_black;
}

// 检查是否有连续的红色节点
bool hasConsecutiveReds(Node* node) {
    if (node == NULL) return false;
    
    // 检查左子节点
    if (node->left != NULL && node->is_red && node->left->is_red) {
        return true;
    }
    
    // 检查右子节点
    if (node->right != NULL && node->is_red && node->right->is_red) {
        return true;
    }
    
    // 递归检查左右子树
    return hasConsecutiveReds(node->left) || hasConsecutiveReds(node->right);
}

// 检查是否为红黑树
bool isRedBlackTree(Node* root) {
    // 性质1：根节点必须是黑色
    if (root == NULL || root->is_red) {
        return false;
    }
    
    // 性质2：检查是否有连续的红色节点
    if (hasConsecutiveReds(root)) {
        return false;
    }
    
    // 性质3：检查所有路径的黑色节点数是否相同
    if (countBlackHeight(root) == -1) {
        return false;
    }
    
    // 检查是否为有效的二叉查找树
    if (!isBST(root, INT_MIN, INT_MAX)) {
        return false;
    }
    
    return true;
}

// 释放树的内存
void freeTree(Node* root) {
    if (root == NULL) return;
    freeTree(root->left);
    freeTree(root->right);
    free(root);
}

int main() {
    int k;
    scanf("%d", &k);
    
    while (k--) {
        int n;
        scanf("%d", &n);
        
        Node* root = NULL;
        for (int i = 0; i < n; i++) {
            int val;
            scanf("%d", &val);
            // 负数表示红色节点
            root = insert(root, val, val < 0);
        }
        
        // 判断是否为红黑树并输出结果
        if (isRedBlackTree(root)) {
            printf("Yes\n");
        } else {
            printf("No\n");
        }
        
        freeTree(root);
    }
    
    return 0;
}