|
| 1 | +import TreeMap |
| 2 | + |
| 3 | +class RedBlackTree(TreeMap): |
| 4 | + """Sorted map implementation using a red-black tree""" |
| 5 | + |
| 6 | + class _Node(TreeMap._Node): |
| 7 | + """Node class for red-black tree mainrains bit that denotes color""" |
| 8 | + __slots__ = '_red' |
| 9 | + |
| 10 | + def __init__(self, element, parent = None, left = None, right = None): |
| 11 | + super().__init__(element, parent, left, right) |
| 12 | + self._red=True # new node red by default |
| 13 | + |
| 14 | + #### positional-based utility methods ##### |
| 15 | + |
| 16 | + #consider nonexistent child to be trivially black |
| 17 | + def _set_red(self, p): p._node._red= True |
| 18 | + def _set_black(self, p): p._node._red = False |
| 19 | + def _set_color(self, p, make_red): p._node._red = make_red |
| 20 | + def _is_red(self, p): return p is not None and p._node._red |
| 21 | + def _is_red_leaf(self, p): return self._is_red(p) and self.is_leaf(p) |
| 22 | + |
| 23 | + def _get_red_child(self, p ): |
| 24 | + """Retrurn a red child of p (or None if no such child)""" |
| 25 | + for child in (self.left(p), self.right(p)): |
| 26 | + if self._is_red(child): |
| 27 | + return child |
| 28 | + |
| 29 | + return None |
| 30 | + |
| 31 | + #support for insertions |
| 32 | + def _rebalance_insert(self, p): |
| 33 | + self._resolve_red(p) # new node is always red |
| 34 | + |
| 35 | + def _resolve_red(self, p): |
| 36 | + if self.is_root(p): |
| 37 | + self._set_black(p) #make root black |
| 38 | + else: |
| 39 | + parent = self.parent(p) |
| 40 | + if self._is_red(parent): #double red problem |
| 41 | + uncle = self.sibling(parent) |
| 42 | + if not self._is_red(uncle): #Case1 misshappen 4-node |
| 43 | + middle=self._restructure(p) # tri node restructuring |
| 44 | + self._set_black(middle) # fix colors |
| 45 | + self._set_red(self.left(middle)) |
| 46 | + self._set_red(self.right(middle)) |
| 47 | + else: # case 2 overfull 5-node |
| 48 | + grand= self.parent(parent) |
| 49 | + self._set_red(grand) #grand parent becomes red |
| 50 | + self._set_black(self.left(grand)) #its children become black |
| 51 | + self._set_black(self.right(grand)) #its children become black |
| 52 | + self._resolve_red(grand) #recur at red grandparent |
| 53 | + |
| 54 | + # Support for deletions |
| 55 | + def _rebalance_delete(self, p): |
| 56 | + if len(self) == 1: |
| 57 | + self._set_black(self.root(p)) |
| 58 | + elif p is not None: |
| 59 | + n = self.num_children(p) |
| 60 | + if n == 1: #deficit exists unless child is a red leaf |
| 61 | + c = next(self.children(p)) |
| 62 | + if not self._is_red_leaf(c): |
| 63 | + self._fix_deficit(p ,c) |
| 64 | + elif n == 2: # removed black node with red child |
| 65 | + if self._is_red_leaf(self.left(p)): |
| 66 | + self._set_black(self.left(p)) |
| 67 | + else: |
| 68 | + self._set_black(self.right(p)) |
| 69 | + |
| 70 | + def _fix_deficit(self, z, y): |
| 71 | + """Resolve black deficit at z where y is the root of z's heavier subtree""" |
| 72 | + |
| 73 | + if not self._is_red(y): # y is black will apply case 1 or case 2 |
| 74 | + x = self._get_red_child(y) |
| 75 | + if x is not None: #case 1 y is black and has red child x; do transfer |
| 76 | + old_color =self._is_red(z) |
| 77 | + middle = self.restructure(x) |
| 78 | + self._set_color(middle, old_color) #middle gets old color of z |
| 79 | + self._set_black(self.left(middle)) |
| 80 | + self._set_black(self.right(middle)) |
| 81 | + else: #Case 2: y is black but no red children; recolor as fusion |
| 82 | + self._set_red(y) |
| 83 | + |
| 84 | + if self._is_red(z): |
| 85 | + self._set_black(z) #resolves problem |
| 86 | + elif not self.is_root(z): |
| 87 | + self._fix_deficit(self.parent(z), self.sibling(z)) #recur upward |
| 88 | + else: #Case 3 y is red rotate misaligned 3-node and repeat |
| 89 | + self.rotate(y) |
| 90 | + self._set_black(y) |
| 91 | + self._set_red(z) |
| 92 | + if z== self.right(y): |
| 93 | + self._fix_deficit(z, self.left(z)) |
| 94 | + else: |
| 95 | + self._fix_deficit(z, self.right(z)) |
| 96 | + |
0 commit comments