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