def _rebalance(self, node): print('rebalancing... node with value {}, current BF = {}'.format( node.key, node.balance_factor)) bst.recurse_bst(self.root, None) 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)
def splay(self, node): """The primary splay operation. Splaying happens all the way up the tree (until `node` no longer has a parent). Three methods are available when splaying: 1. zig step: (1) / (0) 2. zig-zig step: (2) / (1) / (0) 3. zig-zag step: (0) \ (1) \ (2) """ print('\n') print('Splaying node with value: {}'.format(node.key)) print('\n') recurse_bst(self.root, None) while node.parent: if not self.has_grandparent(node): if node.parent.left_child == node: # 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)
def _update_balance(self, node): print('updating balance...') bst.recurse_bst(self.root, None) # Updates the balance factor for all nodes if necessary bf = node.balance_factor # -1, 0 or 1 are considered balanced. Anything else needs a re-balance. if bf > 1 or bf < -1: self._rebalance(node) return # Adjust balance factor for nodes that do not require re-balancing. if node.parent is not None: # If this node has a parent and is a left child, # then increase its balance factor if node.is_left_child(): node.parent.balance_factor += 1 # If parent and right child, decrease its balance factor elif node.is_right_child(): node.parent.balance_factor -= 1 # Update BF for nodes with non-zero balance factors if node.parent.balance_factor != 0: # Keep updating and checking balance until # re-balance has fixed all nodes. self._update_balance(node.parent)
def rank(self, node): ranked = [] for _node in self: ranked.append((_node, _node.size)) # Sort by size, but leave node reference intact. for data in ranked: _node, size = data if _node.name == node: return size if DEBUG: with Section('Order statistic - Binary Search Tree'): stats_bst = OrderStatisticBST() # Create a reasonable spread out tree for x in range(1, 6): stats_bst.put(x * x, {'num': x * x}) stats_bst.put(x * 20, {'num': x * 20}) bst.recurse_bst(stats_bst.root, None) print(stats_bst[1].data) print(stats_bst[100].data) print_h2('Testing select function') for n in range(10): print(stats_bst.select(n)) print_h2('Testing rank function') for n in range(10): print(stats_bst.rank(n))