RBTree.py

from Tree import NIL, Node, Tree

class RBNode(Node):
    def __init__(self, value, color):
        super().__init__(value)
        self.color = color

    def __str__(self):
        s =  f'RBNode:\n'
        s += f'  color: {self.color}\n'
        s += f'  value: {self.value}\n'
        s += f'  parent: {self.parent.value}\n'
        s += f'  left: {self.left.value}\n'
        s += f'  right: {self.right.value}\n'
        s += f'  leaf: {self.leaf()}'
        return s



class RBTree(Tree):
    def __init__(self, root):
        super().__init__(root)

    def leftRotate(self, x):
        y = x.right
        x.right = y.left
        if not y.left == NIL:
            y.left.parent = x
        y.parent = x.parent
        if x.parent == NIL:
            self.root = y
        elif x == x.parent.left:
            x.parent.left = y
        else:
            x.parent.right = y
        y.left = x
        x.parent = y

    def rightRotate(self, y):
        x = y.parent
        x.left = y.right
        if not y.right == NIL:
            y.right.parent = x
        y.parent = x.parent
        if x.parent == NIL:
            self.root = y
        elif x == x.parent.left:
            x.parent.left = y
        else:
            x.parent.right = y
        y.right= x
        x.parent = y


if __name__ == '__main__':
    node = RBNode(7, 'black')
    print(node)

    tree = RBTree(root=node)
    inserted_node = RBNode(8, 'black')
    tree.insert(inserted_node)
    print(node)
    print(inserted_node)

    tree.leftRotate(node)
    print(node)