Last active
September 30, 2022 06:20
-
-
Save BarclayII/cc76c890bb34c74431e40881104fd537 to your computer and use it in GitHub Desktop.
RBTree
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| class Node(object): | |
| def __init__(self, key, red): | |
| self.key = key | |
| self.left = self.right = self.parent = None | |
| self.red = red | |
| @property | |
| def grandparent(self): | |
| if self.parent is None: | |
| return None | |
| return self.parent.parent | |
| @property | |
| def sibling(self): | |
| if self.parent is None: | |
| return None | |
| parent = self.parent | |
| if self is parent.left: | |
| return parent.right | |
| else: | |
| return parent.left | |
| @property | |
| def uncle(self): | |
| if self.parent is None: | |
| return None | |
| return self.parent.sibling | |
| class RBTree(object): | |
| def __init__(self): | |
| self.root = None | |
| def dump(self): | |
| return self._dump(self.root) | |
| def _dump(self, node): | |
| if node is None: | |
| return '' | |
| if node.left is not None: | |
| assert node.left.parent is node | |
| if node.right is not None: | |
| assert node.right.parent is node | |
| return '( %s %s:%s %s )' % ( | |
| self._dump(node.left), | |
| node.key, | |
| 'red' if node.red else 'black', | |
| self._dump(node.right), | |
| ) | |
| def insert(self, key): | |
| node = Node(key, True) | |
| if self.root is None: | |
| self.root = node | |
| return | |
| cur = self.root | |
| while True: | |
| if cur.key == key: | |
| return | |
| elif cur.key < key: | |
| if cur.right is None: | |
| cur.right = node | |
| node.parent = cur | |
| break | |
| else: | |
| cur = cur.right | |
| else: | |
| if cur.left is None: | |
| cur.left = node | |
| node.parent = cur | |
| break | |
| else: | |
| cur = cur.left | |
| self._insert_repair(node) | |
| self.root = node | |
| while self.root.parent is not None: | |
| self.root = self.root.parent | |
| def _rol(self, node): | |
| assert node.right is not None | |
| parent = node.parent | |
| child = node.right | |
| node.right = child.left | |
| if child.left is not None: | |
| child.left.parent = node | |
| child.left = node | |
| node.parent = child | |
| if parent is None: | |
| child.parent = None | |
| elif node is parent.left: | |
| parent.left = child | |
| child.parent = parent | |
| else: | |
| parent.right = child | |
| child.parent = parent | |
| def _ror(self, node): | |
| assert node.left is not None | |
| parent = node.parent | |
| child = node.left | |
| node.left = child.right | |
| if child.right is not None: | |
| child.right.parent = node | |
| child.right = node | |
| node.parent = child | |
| if parent is None: | |
| child.parent = None | |
| elif node is parent.right: | |
| parent.right = child | |
| child.parent = parent | |
| else: | |
| parent.left = child | |
| child.parent = parent | |
| def _insert_repair(self, node): | |
| cur = node | |
| while True: | |
| if cur.parent is None: | |
| # root | |
| break | |
| elif not cur.parent.red: | |
| # parent is black; adding a red node is fine | |
| break | |
| # parent is red... | |
| elif cur.grandparent is None: | |
| # children of the root; simply repaint the root node to black | |
| cur.parent.red = False | |
| break | |
| # parent is red so grandparent must be black... | |
| elif cur.uncle is not None and cur.uncle.red: | |
| # uncle is red; simply repaint the nodes | |
| cur.parent.red = cur.uncle.red = False | |
| cur.grandparent.red = True | |
| cur = cur.grandparent | |
| # parent is red and uncle is black; | |
| else: | |
| if (cur.grandparent.left is not None and | |
| cur is cur.grandparent.left.right): | |
| self._rol(cur.parent) | |
| cur = cur.left | |
| elif (cur.grandparent.right is not None and | |
| cur is cur.grandparent.right.left): | |
| self._ror(cur.parent) | |
| cur = cur.right | |
| if (cur.grandparent.left is not None and | |
| cur is cur.grandparent.left.left): | |
| cur.grandparent.red = True | |
| cur.parent.red = False | |
| self._ror(cur.grandparent) | |
| elif (cur.grandparent.right is not None and | |
| cur is cur.grandparent.right.right): | |
| cur.grandparent.red = True | |
| cur.parent.red = False | |
| self._rol(cur.grandparent) | |
| break | |
| def _remove_repair(self, node): | |
| while node.parent is not None: | |
| if node.sibling is not None and node.sibling.red: | |
| node.parent.red = True | |
| node.sibling.red = False | |
| if node is node.parent.left: | |
| self._rol(node.parent) | |
| else: | |
| self._ror(node.parent) | |
| if not node.parent.red and not node.sibling.red and \ | |
| (node.sibling.left is None or not node.sibling.left.red) and \ | |
| (node.sibling.right is None or not node.sibling.right.red): | |
| node.sibling.red = True | |
| node = node.parent | |
| continue | |
| elif node.parent.red and not node.sibling.red and \ | |
| (node.sibling.left is None or not node.sibling.left.red) and \ | |
| (node.sibling.right is None or not node.sibling.right.red): | |
| node.sibling.red = True | |
| node.parent.red = False | |
| else: | |
| assert node.sibling is None or not node.sibling.red | |
| if node is node.parent.left and (node.sibling.right is None or not node.sibling.right.red): | |
| assert node.sibling.left.red | |
| node.sibling.red = True | |
| node.sibling.left.red = False | |
| self._ror(node.sibling) | |
| elif node is node.parent.right and (node.sibling.left is None or not node.sibling.left.red): | |
| assert node.sibling.right.red | |
| node.sibling.red = True | |
| node.sibling.right.red = False | |
| self._rol(node.sibling) | |
| node.sibling.red = node.parent.red | |
| node.parent.red = False | |
| if node is node.parent.left: | |
| node.right.red = False | |
| self._rol(node.parent) | |
| else: | |
| node.left.red = False | |
| self._ror(node.parent) | |
| break | |
| def remove(self, key): | |
| cur = self.root | |
| while cur is not None: | |
| if cur.key == key: | |
| break | |
| elif cur.key < key: | |
| cur = cur.right | |
| else: | |
| cur = cur.left | |
| if cur is None: | |
| return | |
| # the node to be removed | |
| if cur.left is not None: | |
| node = cur.left | |
| while node.right is not None: | |
| node = node.right | |
| child = node.left | |
| elif cur.right is not None: | |
| node = cur.right | |
| while node.left is not None: | |
| node = node.left | |
| child = node.right | |
| else: | |
| node = cur | |
| child = None | |
| cur.key = node.key | |
| if not node.red and child is None: | |
| self._remove_repair(node) | |
| if child is not None: | |
| assert child.red | |
| child.red = False | |
| p = node.parent | |
| if p is None: | |
| self.root = child | |
| elif p.left is node: | |
| p.left = child | |
| elif p.right is node: | |
| p.right = child | |
| if child is not None: | |
| child.parent = p |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment