算法导论程序31--红黑树的插入（Python）

红黑树的插入过程：

插入过程：rb_insert()

辅助过程：rb_insert_fixup():对结点进行重新着色并旋转

插入结点：z

1.如果z的父结点z.p是黑色的，则直接执行步骤4

2.如果z的父结点是红色的：

（1）z的父结点位于左边的位置：

a.   z的叔结点的颜色是红色：则，设置z的父结点、叔结点为黑色。z的祖父结点为红色。然后z赋值为其祖父结点。

b.   z的叔结点的颜色是黑色：则，分两种情况，z是左孩子还是右孩子

       当z是右孩子时，步骤I：z赋值为z的父结点。左旋z。执行步骤II。

        当z是左孩子时，不必执行步骤I。直接执行步骤II

      II：z的父结点的颜色设置为黑色。z的祖父结点设置为红色。右旋z的祖父结点。

（2）z的父结点位于右边的位置：

a.   z的叔结点的颜色是红色：则，设置z的父结点、叔结点为黑色。z的祖父结点为红色。然后z赋值为其祖父结点。

b.   z的叔结点的颜色是黑色：则，分两种情况，z是左孩子还是右孩子

       当z是左孩子时，步骤I：z赋值为z的父结点。右旋z。执行步骤II。

        当z是右孩子时，不必执行步骤I。直接执行步骤II

      II：z的父结点的颜色设置为黑色。z的祖父结点设置为红色。左旋z的祖父结点。

3.判断z.p的color，如果为红色，继续执行步骤2.

4.设置根结点root.color="BLACK"

class Node:
    def __init__(self,key,right,left,p,color):
        self.key=key
        self.right=right
        self.left=left
        self.p=p
        self.color=color

class tree:
    def __init__(self,root,nil):
        self.root=root
        self.nil=nil
    def tree_insert(self,z):
        y=self.nil
        x=self.root
        while x!=self.nil:
            y=x
            if z.key<x.key:
                x=x.left
            else:
                x=x.right
        z.p=y
        if y==self.nil:
            self.root=z
        elif z.key<y.key:
            y.left=z
        else:
            y.right=z
        z.left=self.nil
        z.right=self.nil
        z.color="RED"
        self.rb_insert_fixup(z)
    def left_rotate(self,x):
        y=x.right
        x.right=y.left
        if y.left!=self.nil:
            y.left.p=x
        y.p=x.p
        if x.p==self.nil:
            self.root=y
        elif x==x.p.left:
            x.p.left=y
        else:
            x.p.right=y
        y.left=x
        x.p=y
        

    def right_rotate(self,y):
        x=y.left
        y.left=x.right
        if x.right!=self.nil:
            x.right.p=y
        x.p=y.p  
        if y.p==self.nil:
            self.root=x
        elif y==y.p.left:
            y.p.left=x
        else:
            y.p.right=x
        x.right=y
        y.p=x

    def rb_insert_fixup(self,z):
        while z.p.color=="RED":
            if z.p==z.p.p.left:
                y=z.p.p.right
                if y.color=="RED":
                    z.p.color="BLACK"
                    y.color="BLACK"
                    z.p.p.color="RED"
                    z=z.p.p
                else:
                    if z==z.p.right:
                        z=z.p
                        self.left_rotate(z)
                    z.p.color="BLACK"
                    z.p.p.color="RED"
                    self.right_rotate(z.p.p)
            else:
                y=z.p.p.left
                if y.color=="RED":
                    z.p.color="BLACK"
                    y.color="BLACK"
                    z.p.p.color="RED"
                    z=z.p.p
                else:
                    if z==z.p.left:
                        z=z.p
                        self.right_rotate(z)
                    z.p.color="BLACK"
                    z.p.p.color="RED"
                    self.left_rotate(z.p.p)
        self.root.color="BLACK"
     def inorder_tree_walk(self,x):
        if x!=self.nil:
            self.inorder_tree_walk(x.left)
            print(x.key)
            self.inorder_tree_walk(x.right)
    def tree_search(self,x,k):
        if x==self.nil or k==x.key:
            return x
        if k < x.key:
            return self.tree_search(x.left,k)
        else: return self.tree_search(x.right,k)
    

运行：

>>> nil=Node(0,None,None,None,"BLACK")
>>> root=Node(7,nil,nil,nil,"BLACK")
>>> t=tree(root,nil)
>>> T=[4,18,3,6,11,19,2,9,14,22,12,17,20]
>>> for i in T:
	z=Node(i,None,None,None,"RED")
	t.tree_insert(z)

	
>>> TT=[7,4,18,3,6,11,19,2,9,14,22,12,17,20]
>>> for i in TT:
	zz=t.tree_search(t.root,i)
	print(i,' ',zz.color)

	
7   BLACK
4   BLACK
18   BLACK
3   BLACK
6   BLACK
11   RED
19   RED
2   RED
9   BLACK
14   BLACK
22   RED
12   RED
17   RED
20   BLACK