[算法导论]红黑树插入实现 @ Python

class RBTree:
    def __init__(self):
        self.nil = RBTreeNode(0)
        self.root = self.nil

class RBTreeNode:
    def __init__(self, x):
        self.key = x
        self.left = None
        self.right = None
        self.parent = None
        self.color = ‘black‘

class Solution:
    def InorderTreeWalk(self, x):
        if x != None:
            self.InorderTreeWalk(x.left)
            if x.key != 0:
                print ‘key:‘, x.key, ‘parent:‘, x.parent.key, ‘color:‘, x.color
            self.InorderTreeWalk(x.right)

    def LeftRotate(self, T, x):
        y = x.right
        x.right = y.left
        if y.left != T.nil:
            y.left.parent = x
        y.parent = x.parent
        if x.parent == T.nil:
            T.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, T, x):
        y = x.left
        x.left = y.right
        if y.right != T.nil:
            y.right.parent = x
        y.parent = x.parent
        if x.parent == T.nil:
            T.root = y
        elif x == x.parent.right:
            x.parent.right = y
        else:
            x.parent.left = y
        y.right = x
        x.parent = y

    def RBInsert(self, T, z):
        # init z
        z.left = T.nil
        z.right = T.nil
        z.parent = T.nil

        y = T.nil
        x = T.root
        while x != T.nil:
            y = x
            if z.key < x.key:
                x = x.left
            else:
                x = x.right
        z.parent = y
        if y == T.nil:
            T.root = z
        elif z.key < y.key:
            y.left = z
        else:
            y.right = z
        z.left = T.nil
        z.right = T.nil
        z.color = ‘red‘
        self.RBInsertFixup(T,z)

    def RBInsertFixup(self, T, z):
        while z.parent.color == ‘red‘:
            if z.parent == z.parent.parent.left:
                y = z.parent.parent.right
                if y.color == ‘red‘:
                    z.parent.color = ‘black‘
                    y.color = ‘black‘
                    z.parent.parent.color = ‘red‘
                    z = z.parent.parent
                else:
                    if z == z.parent.right:
                        z = z.parent
                        self.LeftRotate(T, z)
                    z.parent.color = ‘black‘
                    z.parent.parent.color = ‘red‘
                    self.RightRotate(T,z.parent.parent)
            else:
                y = z.parent.parent.left
                if y.color == ‘red‘:
                    z.parent.color = ‘black‘
                    y.color = ‘black‘
                    z.parent.parent.color = ‘red‘
                    z = z.parent.parent
                else:
                    if z == z.parent.left:
                        z = z.parent
                        self.RightRotate(T, z)
                    z.parent.color = ‘black‘
                    z.parent.parent.color = ‘red‘
                    self.LeftRotate(T, z.parent.parent)
        T.root.color = ‘black‘

nodes = [11,2,14,1,7,15,5,8,4]
T = RBTree()
s = Solution()
for node in nodes:
    s.RBInsert(T,RBTreeNode(node))

s.InorderTreeWalk(T.root)