From dca7598c5176861c2ca3f827991a76161015b57e Mon Sep 17 00:00:00 2001
From: jf-home <Orig5826@163.com>
Date: Sun, 23 Jun 2024 23:40:27 +0800
Subject: [PATCH] =?UTF-8?q?avl=E6=A0=91=E8=B0=83=E8=AF=95=E9=80=9A?=
 =?UTF-8?q?=E8=BF=87=EF=BC=8C=E5=BE=88nice=EF=BC=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/tree.c       | 81 ++++++++++++++++++++++++++++++------------------
 test/tree_test.c |  2 +-
 2 files changed, 51 insertions(+), 32 deletions(-)

diff --git a/src/tree.c b/src/tree.c
index de323d6..ccec59f 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -1402,6 +1402,16 @@ static bool tree_rebalance(ptree_t head, ptree_node_t tree)
 
 #endif
 
+static void tree_set_balance(struct _tree* self, struct _tree_node * node)
+{
+    assert(self != NULL);
+    if(node == NULL)
+    {
+        return;
+    }
+    node->balance = self->height(self, node->right) - self->height(self, node->left);
+}
+
 static struct _tree_node* tree_turn_left(struct _tree* self, struct _tree_node* root)
 {
     assert(self != NULL);
@@ -1409,14 +1419,25 @@ static struct _tree_node* tree_turn_left(struct _tree* self, struct _tree_node*
     assert(root->right != NULL);
     struct _tree_node* node = root->right;
     
-    root->parent->right = node;         // step1
+    if(root->parent != NULL)
+    {
+        if(root->parent->left == root)          // step1
+        {
+            root->parent->left = node;
+        }
+        else if(root->parent->right == root)
+        {
+            root->parent->right = node;         // setp1
+        }
+    }
     node->parent = root->parent;        // step2
-
     root->parent = node;                // step3
-    node->left = root;                  // step4
+
+    root->right = node->left;           // step4
+    node->left = root;                  // step5
     
-    root->balance = 0;
-    node->balance = 0;
+    tree_set_balance(self, root);
+    tree_set_balance(self, node);
     return node;
 }
 
@@ -1427,13 +1448,25 @@ static struct _tree_node* tree_turn_right(struct _tree* self, struct _tree_node*
     assert(root->left != NULL);
     struct _tree_node* node = root->left;
     
-    root->parent->left = node;          // step1
+    if(root->parent != NULL)
+    {
+        if(root->parent->left == root)
+        {
+            root->parent->left = node;          // step1
+        }
+        else if(root->parent->right == root)
+        {
+            root->parent->right = node;         // setp1
+        }
+    }
     node->parent = root->parent;        // step2
     root->parent = node;                // step3
-    node->right = root;                 // step4
+
+    root->left = node->right;           // step4
+    node->right = root;                 // step5
     
-    root->balance = 0;
-    node->balance = 0;
+    tree_set_balance(self, root);
+    tree_set_balance(self, node);
     return node;
 }
 
@@ -1467,26 +1500,9 @@ uint32_t tree_height(struct _tree* self, struct _tree_node* root)
     {
         return 0;
     }
-    uint32_t left = tree_height(self, root->left);
-    uint32_t right = tree_height(self, root->right);
-    if(left > right)
-    {
-        return left + 1;
-    }
-    else
-    {
-        return right + 1;
-    }
-}
-
-static void tree_set_balance(struct _tree* self, struct _tree_node * node)
-{
-    assert(self != NULL);
-    if(node == NULL)
-    {
-        return;
-    }
-    node->balance = self->height(self, node->right) - self->height(self, node->left);
+    uint32_t left_height = tree_height(self, root->left);
+    uint32_t right_height = tree_height(self, root->right);
+    return (left_height > right_height) ? (left_height + 1) : (right_height + 1);
 }
 
 /**
@@ -1537,11 +1553,15 @@ static bool tree_avl_rebalance(struct _tree* self, struct _tree_node* root)
             root = tree_trun_right_then_left(self, root);
         }
     }
-
+    
     if(root->parent != NULL)
     {
         tree_avl_rebalance(self, root->parent);
     }
+    else
+    {
+        self->_root = root;
+    }
     return true;
 }
 
@@ -1671,7 +1691,6 @@ bool tree_avl_insert(struct _tree* self, void* obj)
         }
         self->rebalance(self, root);
     }
-
     self->_size++;
     return true;
 }
diff --git a/test/tree_test.c b/test/tree_test.c
index 74142f8..78014c6 100644
--- a/test/tree_test.c
+++ b/test/tree_test.c
@@ -279,7 +279,7 @@ void tree_test(void)
 void test_tree_num(void)
 {
     uint32_t i = 0;
-    int data[] = { 1,2,3,4,5,6,7,8,9,10 };
+    int data[] = { 1,2,3,4,5,6};
     // int data[] = { 5,2,3,1,7,8,6 };
     int temp = 0;
     uint32_t len = sizeof(data) / sizeof(data[0]);