if node.balance_factor < 0: # If the current nodes balanace factor is < 0, and its right # child's balance factor is > 0, rotate both sides, right and left. # Otherwise, just rotate the left side, if the right child is # not unbalanced. if node.right_child is not None: if node.right_child.balance_factor > 0: self._rotate_right(node.right_child) self._rotate_left(node) else: self._rotate_left(node) # If the balance factor of this node is > 0, then do the opposite # rotation for both child nodes. elif node.balance_factor > 0: if node.left_child is not None: if node.left_child.balance_factor < 0: self._rotate_left(node.left_child) self._rotate_right(node) else: self._rotate_right(node) print('New BF for node with value {} is {}'.format( node.key, node.balance_factor)) bst.recurse_bst(self.root, None) if __name__ == '__main__': with Section('AVL Trees'): avl = AVLTree() bst.populate_bst(avl, count=5)
# Zig right self.right_rotate(node.parent) else: # Zig left self.left_rotate(node.parent) elif node.parent.left_child == node and \ node.parent.parent.left_child == node.parent: # Zig-zig right self.right_rotate(node.parent.parent) self.right_rotate(node.parent) elif node.parent.right_child == node and \ node.parent.parent.right_child == node.parent: # Zig-zig left self.left_rotate(node.parent.parent) self.left_rotate(node.parent) elif node.parent.left_child == node and \ node.parent.parent.right_child == node.parent: # Zig-zag right self.right_rotate(node.parent) self.left_rotate(node.parent) else: # Zig-zag left self.left_rotate(node.parent) self.right_rotate(node.parent) if __name__ == '__main__': with Section('Splay Tree'): splay = SplayTree() populate_bst(splay, count=10)