From 68706331a1694a16d512990c5117b293ac0c4b40 Mon Sep 17 00:00:00 2001
From: wjf-hs <Orig5826@163.com>
Date: Tue, 2 Jul 2024 14:43:12 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=9E=E7=8E=B0tree=5Frb=5Fcheck=E6=96=B9?=
 =?UTF-8?q?=E4=BE=BF=E7=9B=B4=E6=8E=A5=E5=AE=9A=E4=BD=8D=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/tree.c       |   2 +-
 test/test_tree.c | 125 ++++++++++++++++++-----------------------------
 2 files changed, 49 insertions(+), 78 deletions(-)

diff --git a/src/tree.c b/src/tree.c
index 94f6d3c..7ea77aa 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -1470,7 +1470,7 @@ bool tree_rb_delete(struct _tree* self, void* obj)
     if(tmp->color == RBT_BLACK)
     {
         tree_rb_delete_fix(self, tmp);
-    }         
+    }
 
     if(tmp->parent != NULL)
     {
diff --git a/test/test_tree.c b/test/test_tree.c
index c478005..cde98e2 100644
--- a/test/test_tree.c
+++ b/test/test_tree.c
@@ -126,85 +126,49 @@ void test_avltree_num(void)
     tree_free(tree);
 }
 
-bool tree_rb_check(struct _tree* self, struct _tree_node* root)
+static bool tree_rb_check_color(struct _tree *self, struct _tree_node* root, int black_num, int black_num_expected)
 {
-    assert(self != NULL);
     if(root == NULL)
+    {
+        if(black_num != black_num_expected)
+        {
+            return false;
+        }
+        return true;
+    }
+    if(root->color == RBT_BLACK)
+    {
+        black_num++;
+    }
+    if(root->color == RBT_RED && root->parent && root->parent->color == RBT_RED)
     {
         return false;
     }
+    return tree_rb_check_color(self, root->left, black_num, black_num_expected) && 
+        tree_rb_check_color(self, root->right, black_num, black_num_expected);
+}
 
-    struct _tree_node* node = root;
-    queue_t queue = queue_new();
-    queue_init(queue, sizeof(struct _tree_node*));
+static bool tree_rb_check(struct _tree* self)
+{
+    assert(self != NULL);
 
-    if(node != NULL)
+    if(self->_root->color != RBT_BLACK)
     {
-        queue->push(queue, &node);
-        while(!queue->empty(queue))
-        {
-            queue->pop(queue, &node);
-            if(!self->_right_priority)
-            {
-                if(node->left != NULL)
-                {
-                    queue->push(queue, &node->left);
-                }
-                if(node->right != NULL)
-                {
-                    queue->push(queue, &node->right);
-                }
-            }
-            else
-            {
-                if(node->right != NULL)
-                {
-                    queue->push(queue, &node->right);
-                }
-                if(node->left != NULL)
-                {
-                    queue->push(queue, &node->left);
-                }
-            }
-            
-            if(node->color == RBT_RED)
-            {
-                if(node->parent != NULL)
-                {
-                    if(node->parent->color == RBT_RED)
-                    {
-                        return false;
-                    }
-                }
-
-                if(node->left != NULL)
-                {
-                    if(node->left->color == RBT_RED)
-                    {
-                        return false;
-                    }
-                }
-
-                if(node->right != NULL)
-                {
-                    if(node->right->color == RBT_RED)
-                    {
-                        return false;
-                    }
-                }
-            }
-
-            if(node->parent == NULL)
-            {
-                if(self->_root != node || node->color != RBT_BLACK)
-                {
-                    return false;
-                }
-            }
-        }
+        printf("self->_root->color != RBT_BLACK\n");
+        return false;
     }
-    queue_free(queue);
-    return true;
+
+    int black_num_expected = 0;
+    struct _tree_node* root = self->_root;
+    while(root)
+    {
+        if(root->color == RBT_BLACK)
+        {
+            black_num_expected++;
+        }
+        root = root->left;
+    }
+    return tree_rb_check_color(self, self->_root, 0, black_num_expected);
 }
 
 /**
@@ -238,14 +202,12 @@ void test_rbtree_num(void)
         printf("size = %2d : ", tree->size(tree));
         tree->preorder(tree, tree->_root);
         printf("\n");
-    }
 
-    printf("----- rb_check -----\n");
-    bool ret = tree_rb_check(tree, tree->_root);
-    if(ret != true)
-    {
-        printf("----- rb_check_error -----\n");
-        return;
+        if(true != tree_rb_check(tree))
+        {
+            printf("----- rb_check_error -----\n");
+            return;
+        }
     }
 
     printf("----- max -----\n");
@@ -316,6 +278,9 @@ void test_rbtree_num(void)
     for (i = 0; i < len; i++)
     {
         temp = data[i];
+        printf("delete = ");
+        tree->print_obj(&temp);
+
         // delete
         tree->delete(tree, &temp);
 
@@ -324,6 +289,12 @@ void test_rbtree_num(void)
         printf("size = %2d : ", tree->size(tree));
         tree->preorder(tree, tree->_root);
         printf("\n");
+
+        if(true != tree_rb_check(tree))
+        {
+            printf("----- rb_check_error -----\n");
+            return;
+        }
     }
 
     if(tree->empty(tree))