算法导论程序32--红黑树的删除

红黑树的删除

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)

    def rb_transplant(self,u,v):
        if u.p==self.nil:
            self.root=v
        elif u==u.p.left:
            u.p.left=v
        else:
            u.p.right=v
        v.p=u.p
    def tree_minimum(self,x):
        while x.left!=self.nil:
            x=x.left
        return x

    
    def rb_delete(self,z):
        y=z
        y_original_color=y.color
        if z.left==self.nil:
            x=z.right
            self.rb_transplant(z,z.right)
        elif z.right==self.nil:
            x=z.left
            self.rb_transplant(z,z.left)
        else:
            y=self.tree_minimum(z.right)
            y_original_color=y.color
            x=y.right
            if y.p==z:
                x.p=y
            else:
                self.rb_transplant(y,y.right)
                y.right=z.right
                y.right.p=y
            self.rb_transplant(z,y)
            y.left=z.left
            y.left.p=y
            y.color=z.color
            if y_original_color=="BLACK":
                self.rb_delete_fixup(x)    
            
            

    def rb_delete_fixup(self,x):
        while x!=self.root and x.color=="BLACK":
            if x==x.p.left:
                w=x.p.right
                if w.color=="RED":
                    w.color="BLACK"
                    x.p.color="RED"
                    self.left_rotate(x.p)
                    w=x.p.right
                if w.left.color=="BLACK" and w.right.color=="BLACK":
                    w.color="RED"
                    x=x.p
                else:
                    if w.right.color=="BLACK":
                        w.left.color=="BLACK"
                        w.color="RED"
                        self.right_rotate(w)
                        w=x.p.right
                    w.color=x.p.color
                    x.p.color="BLACK"
                    w.right.color="BLACK"
                    self.left_rotate(x.p)
                    x=self.root
            else:
                w=x.p.left
                if w.color=="RED":
                    w.color="BLACK"
                    x.p.color="RED"
                    self.right_rotate(x.p)
                    w=x.p.left
                if w.right.color=="BLACK" and w.left.color=="BLACK":
                    w.color="RED"
                    x=x.p
                else:
                    if w.left.color=="BLACK":
                        w.right.color=="BLACK"
                        w.color="RED"
                        self.left_rotate(w)
                        w=x.p.left
                    w.color=x.p.color
                    x.p.color="BLACK"
                    w.left.color="BLACK"
                    self.right_rotate(x.p)
                    x=self.root
        x.color="BLACK"
                    
 
    

运行结果：

>>> 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,nil,nil,nil,"RED")
	t.tree_insert(z)

	
>>> t.rb_delete(t.tree_search(t.root,4))

>>> t.root.left.key
3
>>> t.root.left.left.key
2
>>> t.root.left.right.key
6

>>> t.root.left.color
'BLACK'
>>> t.root.left.left.color
'BLACK'
>>> t.root.left.right.color
'BLACK'