Python之红黑树

NIL = Node(-1, 1, None, None)
NIL.left = NIL
NIL.right = NIL
class Node:
    def __init__(self, key = -1, color = 0, left = NIL, right = NIL):
        self.key = key
        self.color = color
        self.left = left
        self.right = right
    def __repr__(self):
        return "[" + str(self.key) + "-" + str(self.color) + ": (" + str(self.left.key) + "-" + str(self.left.color) + ", " + str(self.right.key) + "-" + str(self.right.key) + ")]"     

class RBTree:
    def __init__(self, root = NIL):
        self.root = root
    def insert(self, key):
        self.root = self.__insert(self.root, key)
        self.root.color = 1
        return self.root
    def __insert(self, root, key):
        if root == NIL:
            return Node(key)
        if root.key == key:
            return root
        if root.key < key:
            root.right = self.__insert(root.right, key)
        else:
            root.left = self.__insert(root.left, key)
        return self.insert_maintain(root)
    def insert_maintain(self, root):
        if not self.has_red_child(root):
            return root
        if root.left.color == 0 and root.right.color == 0:
            if not self.has_red_child(root.left) and not self.has_red_child(root.right):
                return root
            root.color = 0
            root.left.color = root.right.color = 1
            return root
        if root.left.color == 0 and not self.has_red_child(root.left):
            return root
        if root.right.color == 0 and not self.has_red_child(root.right):
            return root
        if root.left.color == 0:
            if root.left.right.color == 0:
                root.left = self.left_rotate(root.left)
            root = self.right_rotate(root)
        else:
            if root.right.left.color == 0:
                root.right = self.right_rotate(root.right)
            root = self.left_rotate(root)
        root.color = 0
        root.left.color = root.right.color = 1
        return root
    def erase(self, key):
        self.root = self.__erase(self.root, key)
        self.root.color = 1
        return self.root
    def __erase(self, root, key):
        if root == NIL:
            return NIL
        if root.key < key:
            root.right = self.__erase(root.right, key)
        elif root.key > key:
            root.left = self.__erase(root.left, key)
        else:
            if root.left == NIL or root.right == NIL:
                tmp = root.left if root.left != NIL else root.right
                tmp.color += root.color
                del root
                return tmp
            tmp = self.get_pre(root)
            root.key = tmp.key
            root.left = self.__erase(root.left, tmp.key)
        return self.erase_maintain(root)
    def erase_maintain(self, root):
        if root.left.color != 2 and root.right.color != 2:
            return root
        if self.has_red_child(root):
            root.color = 0
            if root.left.color == 0:
                root = self.right_rotate(root)
                root.right = self.erase_maintain(root.right)
            else:
                root = self.left_rotate(root)
                root.left = self.erase_maintain(root.left)
            root.color = 1
            return root
        if root.left.color == 2 and not self.has_red_child(root.right) or root.right.color == 2 and not self.has_red_child(root.left):
            root.left.color -= 1
            root.right.color -= 1
            root.color += 1
            return root
        if root.left.color == 2:
            root.left.color = 1
            if root.right.right.color != 0:
                root.right.color = 0
                root.right = self.right_rotate(root.right)
                root.right.color = 1
            root = self.left_rotate(root)
            root.color = root.left.color
        else:
            root.right.color = 1
            if root.left.left.color != 0:
                root.left.color = 0
                root.left = self.left_rotate(root.left)
                root.left.color = 1
            root = self.right_rotate(root)
            root.color = root.right.color
        root.left.color = root.right.color = 1
        return root
    def left_rotate(self, root):
        tmp = root.right
        root.right = tmp.left
        tmp.left = root
        return tmp
    def right_rotate(self , root):
        tmp = root.left
        root.left = tmp.right
        tmp.right = root
        return tmp
    def get_pre(self, root):
        tmp = root.left
        while tmp.right != NIL:
            tmp = tmp.right
        return tmp
    def has_red_child(self, root):
        return root.left.color == 0 or root.right.color == 0
    def inorder(self, root):
        if root.left != NIL:
            self.inorder(root.left)
        print(root)
        if root.right != NIL:
            self.inorder(root.right)
        return None

t = RBTree()
for i in range(1, 7):
    t.insert(i)
    if i % 3 == 0:
        t.erase(i - 1)
t.inorder(t.root)