correctly handle rebalance on uneven depth (#47)

CKS-Systems · Aug 29, 2024 · f1a87ca · f1a87ca
1 parent 2038f3c
commit f1a87ca
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/lib/src/red_black_tree.rs b/lib/src/red_black_tree.rs
@@ -688,6 +688,11 @@ impl<'a, V: TreeValue> RedBlackTree<'a, V> {
 
         trace!("FIX G=[{grandparent_index}:{grandparent_color:?}] P=[{parent_index}:{parent_color:?}] Pi={parent_is_left} Ci={current_is_left}");
 
+        if grandparent_index == NIL && parent_color == Color::Red {
+            self.set_color(parent_index, Color::Black);
+            return;
+        }
+
         // Case II: Uncle is black, left left
         if parent_is_left && current_is_left {
             self.rotate_right(grandparent_index);
@@ -698,11 +703,9 @@ impl<'a, V: TreeValue> RedBlackTree<'a, V> {
         // Case III: Uncle is black, left right
         if parent_is_left && !current_is_left {
             self.rotate_left(parent_index);
+            self.rotate_right(grandparent_index);
             self.set_color(index_to_fix, grandparent_color);
-            if grandparent_index != NIL {
-                self.rotate_right(grandparent_index);
-                self.set_color(grandparent_index, index_to_fix_color);
-            }
+            self.set_color(grandparent_index, index_to_fix_color);
         }
         // Case IV: Uncle is black, right right
         if !parent_is_left && !current_is_left {
@@ -713,11 +716,10 @@ impl<'a, V: TreeValue> RedBlackTree<'a, V> {
         // Case V: Uncle is black, right left
         if !parent_is_left && current_is_left {
             self.rotate_right(parent_index);
+            self.rotate_left(grandparent_index);
             self.set_color(index_to_fix, grandparent_color);
-            if grandparent_index != NIL {
-                self.rotate_left(grandparent_index);
-                self.set_color(grandparent_index, index_to_fix_color);
-            }
+            self.set_color(grandparent_index, index_to_fix_color);
+
         }
     }