def update_balance(self, node: BinaryTreeNode): if node.balance_factor > 1 or node.balance_factor < -1: self.rebalance(node) return if node.has_parent(): if node.is_left_child(): node.parent.balance_factor += 1 elif node.is_right_child(): node.parent.balance_factor -= 1 if node.parent.balance_factor != 0: self.update_balance(node.parent)
def _get(self, node: BinaryTreeNode, key): if key == node.key: return node.value elif key < node.key: if node.has_left_child(): return self._get(node.left, key) else: raise KeyError elif key > node.key: if node.has_right_child(): return self._get(node.right, key) else: raise KeyError
def _delete(self, node: BinaryTreeNode, key): if not node: return None if key < node.key: node.left = self._delete(node.left, key) elif key > node.key: node.right = self._delete(node.right, key) else: if not node.right: return node.left if not node.left: return node.right successor = self._delete_min(node.right) node.parent.set_child(successor) successor.set_left(node.left) successor.set_right(node.right) node = successor node.update_size() return node
def _put(self, node: BinaryTreeNode, key, value): if key == node.key: node.value = value elif key < node.key: if node.has_left_child(): self._put(node.left, key, value) else: node.left = BinaryTreeNode(key, value, parent=node) else: if node.has_right_child(): self._put(node.right, key, value) else: node.right = BinaryTreeNode(key, value, parent=node) node.update_size()
def _put(self, node: BinaryTreeNode, key, value): # same as bst. if key == node.key: node.value = value return elif key < node.key: if node.has_left_child(): self._put(node.left, key, value) else: node.left = BinaryTreeNode(key, value, parent=node) else: if node.has_right_child(): self._put(node.right, key, value) else: node.right = BinaryTreeNode(key, value, parent=node) # rotate and colors change if self.is_red(node.right) and not self.is_red(node.left): node = self.rotate_left(node) if self.is_red(node.left) and node.left and self.is_red(node.left.left): node = self.rotate_right(node) self.flip_colors(node) node.update_size()
def rotate_right(self, old_root: BinaryTreeNode): new_root = old_root.left # 1.change upper relationships. if old_root.is_root(): self.set_root(new_root) else: old_root.parent.set_child(new_root) # 2.change old_root's left old_root.set_left(new_root.right) # 3.change new_root's right new_root.set_right(old_root) # 4.update two node's size old_root.update_size() new_root.update_size() # 5.change balance factor old_root.balance_factor = old_root.balance_factor - 1 - max(0, new_root.balance_factor) new_root.balance_factor = new_root.balance_factor + 1 + max(0, old_root.balance_factor)
def rotate_right(self, old_root: BinaryTreeNode): new_root = old_root.left # 1.change upper relationships. if old_root.is_root(): self.set_root(new_root) else: old_root.parent.set_child(new_root) # 2.change old_root's left old_root.set_left(new_root.right) # 3.change new root's right new_root.set_right(old_root) # 4.update color new_root.color = old_root.color old_root.color = RED # 5.update size old_root.update_size() new_root.update_size() return new_root
def set_root(self, node: BinaryTreeNode): if node: node.parent = None self.root = node
def flip_colors(self, node: BinaryTreeNode): if self.is_red(node.left) and self.is_red(node.right): node.left.color = BLACK node.right.color = BLACK if not node.is_root(): node.color = RED
def is_red(node: BinaryTreeNode): if node: return node.is_red() else: return False