fixed error and added test method

PhoenixIra · Nov 15, 2023 · 6061b70 · 6061b70
1 parent f1ecbf6
commit 6061b70
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/pyalgotask/tasks/trees/avl_tree.py b/pyalgotask/tasks/trees/avl_tree.py
@@ -74,7 +74,8 @@ def insert(self, value):
         self.adjust_depth(new_node.parent)
 
         # rebalance (A7-A9)
-        self.balance(new_node)
+        root = self.balance(new_node)
+        self.algorithm_root = root
 
     def adjust_depth(self, node):
         """adjusts the depth values starting from node
@@ -140,8 +141,9 @@ def balance(self, node):
         :return: the new root of the tree
         """
         if node.is_nil():
-            return
+            return node
 
+        child = node
         current = node.parent
         while not current.is_nil():
             if current.left.depth == current.right.depth + 2:
@@ -151,8 +153,11 @@ def balance(self, node):
                 # right imbalance
                 current = self.balance_right(current)
 
+            child = current
             current = current.parent
 
+        return child
+
     def delete(self, value):
         """Performs a delete of the value from the tree
         :param value: the value to delete
@@ -209,7 +214,7 @@ def algorithm(self):
                 raise NotImplementedError(
                     f"Operator {operation.type} not supported by task {self.__class__.__name__}"
                 )
-            yield (self.algorithm_root, str(operation))
+            yield (self.algorithm_root, operation)
 
 
 task_base.register_task("tree", AVLTree())
diff --git a/pyalgotask/tasks/trees/binary_search_tree.py b/pyalgotask/tasks/trees/binary_search_tree.py
@@ -146,7 +146,7 @@ def algorithm(self):
                 raise NotImplementedError(
                     f"Operator {operation.type} not supported by task {self.__class__.__name__}"
                 )
-            yield (self.algorithm_root, str(operation))
+            yield (self.algorithm_root, operation)
 
 
 task_base.register_task("tree", BinarySearchTree())