feat: issue #3 adds test that takes 4 previously untested branches

Author: ztick

data_structures/binary_tree/red_black_tree.py

@@ -621,6 +621,31 @@ def test_insert() -> bool:
     ans.right.right.right = RedBlackTree(12, 1, ans.right.right)
     return tree == ans
 
+def test_remove() -> bool:
+    """Test the remove() method of the tree correctly balances, colors,
+    and removes. The tested tree should be the same as in test_insert()
+    to rule out any insertion errors.
+    """
+    tree = RedBlackTree(0)
+    tree.insert(8)
+    tree.insert(-8)
+    tree.insert(-5)
+    tree.insert(-6)
+    tree.insert(4)
+    tree.insert(12)
+    tree.insert(10)
+    tree.insert(11)
+    tree.remove(-6)
+    ans = RedBlackTree(0, 0)
+    ans.left = RedBlackTree(-8, 0, ans)
+    ans.left.right = RedBlackTree(-5, 1, ans)
+    ans.right = RedBlackTree(8, 1, ans)
+    ans.right.left = RedBlackTree(4, 0, ans.right)
+    ans.right.right = RedBlackTree(11, 0, ans.right)
+    ans.right.right.left = RedBlackTree(10, 1, ans.right.right)
+    ans.right.right.right = RedBlackTree(12, 1, ans.right.right)
+    return tree == ans
+
 
 def test_insert_and_search() -> bool:
     """Tests searching through the tree for values."""
@@ -748,6 +773,7 @@ def main() -> None:
     """
     print_results("Rotating right and left", test_rotations())
     print_results("Inserting", test_insert())
+    print_results("Removing", test_remove())
     print_results("Searching", test_insert_and_search())
     print_results("Deleting", test_insert_delete())
     print_results("Floor and ceil", test_floor_ceil())