diff --git a/avl.py b/avl.py index 5a6c2b0..00f16b5 100644 --- a/avl.py +++ b/avl.py @@ -124,20 +124,17 @@ class AVLTree(RankedTree[T]): if rotating_around_root: self.root = new_root - return factor == 0 + return factor != 0 def __delete_fixup( - self, y: Optional[Node[T]], parent: Optional[Node[T]] = None + self, x: Optional[Node[T]], parent: Optional[Node[T]] = None ) -> bool: - x = y if y else parent - assert x - factor = _balance_factor(x) if factor == 0: _update_rank(x) - return False - elif factor in (-1, 1): return True + elif factor in (-1, 1): + return False rotating_around_root = x.parent is None y, leaning, to_left, to_right = ( @@ -158,11 +155,13 @@ class AVLTree(RankedTree[T]): def _delete_rebalance( self, node: Optional[Node[T]], parent: Optional[Node[T]] ) -> None: - while node or parent: - # TODO: Check if it is possible to not propagate all the way up. - # if self.__delete_fixup(node, parent): - # return - self.__delete_fixup(node, parent) - node, parent = parent, (parent.parent if parent else None) + if not node and not parent: + return + + if not node: + node, parent = parent, parent.parent + + while node and self.__delete_fixup(node, parent): + node, parent = parent, parent.parent if parent else None # endregion DeleteRebalance