python 实现数据结构 课程设计

""" psf/black : true ruff : passed """ from __future__ import annotations from collections.abc import Iterator class RedBlackTree: """ A Red-Black tree, which is a self-balancing BST (binary search tree). This tree has similar performance to AVL trees, but the balancing is less strict, so it will perform faster for writing/deleting nodes and slower for reading in the average case, though, because they're both balanced binary search trees, both will get the same asymptotic performance. To read more about them, https://en.wikipedia.org/wiki/Red–black_tree Unless otherwise specified, all asymptotic runtimes are specified in terms of the size of the tree. """ def __init__( self, label: int | None = None, color: int = 0, parent: RedBlackTree | None = None, left: RedBlackTree | None = None, right: RedBlackTree | None = None, ) -> None: """Initialize a new Red-Black Tree node with the given values: label: The value associated with this node color: 0 if black, 1 if red parent: The parent to this node left: This node's left child right: This node's right child """ self.label = label self.parent = parent self.left = left self.right = right self.color = color # Here are functions which are specific to red-black trees def rotate_left(self) -> RedBlackTree: """Rotate the subtree rooted at this node to the left and returns the new root to this subtree. Performing one rotation can be done in O(1). """ parent = self.parent right = self.right if right is None: return self self.right = right.left if self.right: self.right.parent = self self.parent = right right.left = self if parent is not None: if parent.left == self: parent.left = right else: parent.right = right right.parent = parent return right def rotate_right(self) -> RedBlackTree: """Rotate the subtree rooted at this node to the right and returns the new root to this subtree. Performing one rotation can be done in O(1). """ if self.left is None: return self parent = self.parent left = self.left self.left = left.right if self.left: self.left.parent = self self.parent = left left.right = self if parent is not None: if parent.right is self: parent.right = left else: parent.left = left left.parent = parent return left def insert(self, label: int) -> RedBlackTree: """Inserts label into the subtree rooted at self, performs any rotations necessary to maintain balance, and then returns the new root to this subtree (likely self). This is guaranteed to run in O(log(n)) time. """ if self.label is None: # Only possible with an empty tree self.label = label return self if self.label == label: return self elif self.label > label: if self.left: self.left.insert(label) else: self.left = RedBlackTree(label, 1, self) self.left._insert_repair() else: if self.right: self.right.insert(label) else: self.right = RedBlackTree(label, 1, self) self.right._insert_repair() return self.parent or self def _insert_repair(self) -> None: """Repair the coloring from inserting into a tree.""" if self.parent is None: # This node is the root, so it just needs to be black self.color = 0 elif color(self.parent) == 0: # If the parent is black, then it just needs to be red self.color = 1 else: uncle = self.parent.sibling if color(uncle) == 0: if self.is_left() and self.parent.is_right(): self.parent.rotate_right() if self.right: self.right._insert_repair() elif self.is_right() and self.parent.is_left(): self.parent.rotate_left() if self.left: self.left._insert_repair() elif self.is_left(): if self.grandparent: self.grandparent.rotate_right() self.parent.color = 0 if self.parent.right: self.parent.right.color = 1 else: if self.grandparent: self.grandparent.rotate_left() self.parent.color = 0 if self.parent.left: self.parent.left.color = 1 else: self.parent.color = 0 if uncle and self.grandparent: uncle.color = 0 self.grandparent.color = 1 self.grandparent._insert_repair() def remove(self, label: int) -> RedBlackTree: """Remove label from this tree.""" if self.label == label: if self.left and self.right: # It's easier to balance a node with at most one child, # so we replace this node with the greatest one less than # it and remove that. value = self.left.get_max() if value is not None: self.label = value self.left.remove(value) else: # This node has at most one non-None child, so we don't # need to replace child = self.left or self.right if self.color == 1: # This node is red, and its child is black # The only way this happens to a node with one child # is if both children are None leaves. # We can just remove this node and call it a day. if self.parent: if self.is_left(): self.parent.left = None else: self.parent.right = None else: # The node is black if child is None: # This node and its child are black if self.parent is None: # The tree is now empty return RedBlackTree(None) else: self._remove_repair() if self.is_left(): self.parent.left = None else: self.parent.right = None self.parent = None else: # This node is black and its child is red # Move the child node here and make it black self.label = child.label self.left = child.left self.right = child.right if self.left: self.left.parent = self if self.right: self.right.parent = self elif self.label is not None and self.label > label: if self.left: self.left.remov