算法导论程序33--动态顺序统计

顺序统计：

n个元素集合中的第i（i属于{1,2,...,n}）个顺序统计量就是简单地规定该集合中的具有第i小关键字的元素。

本文介绍如何修改红黑树，使得可以在O（lgn）时间内确定任何的顺序统计量。

顺序统计树T只是简单地在每个结点上存储附加信息的一棵红黑树，在红黑树的结点x中，除了通常属性x.key，x.color，x.p，x.left和x.right之外，还包括另一个属性x.size。

这个属性包含了以x为根的子树（包括x本身）的（内）结点数，即这课子树的大小。

如果定义哨兵的size大小为0，也就是设置T.nil.size为0。则有等式：

x.size=x.left.size+x.right.size+1

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

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"
        #z插入成功后，size初始化为1
        z.size=1
        #红黑树性质维护
        self.rb_insert_fixup(z)
        #更新结点z的父结点直到根结点的size属性
        while z.p!=self.nil:
            z.p.size+=1
            z=z.p
            
    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
        #左旋导致两个结点的size属性失效，size属性更新如下：
        y.size=x.size
        x.size=x.left.size+x.right.size+1
        

    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
        #右旋导致两个结点的size属性失效，size属性更新如下：
        x.size=y.size
        y.size=y.right.size+y.left.size+1

    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"

    def print_tree(self,z):
        if z!=self.nil:
            print(z.key,z.color,":",end='')
            print("( ",end='')
            print(z.left.key,z.left.color,"   ",end='')
            print(z.right.key,z.right.color,end='')
            print(" )",end='')
            
    #找出统计树中的第i小的关键字
    def os_select(self,x,i):
        r=x.left.size+1
        if i==r:
            return x
        elif x<r:
            return self.os_select(x.left,i)
        else:
            #注意这种情况就应该找以x为根结点的右子树中第i-r小的关键字了。
            return self.os_select(x.right,i-r)
                
    #确定一个元素的秩:返回对T中序遍历对应的线性序中x的位置
    def os_rank(self,x):
        r=x.left.size+1
        y=x
        while y!=self.root:
            if y==y.p.right:
                r+=y.p.left.size+1
            y=y.p
        return r
            
        

    
    

运行结果：

>>> nil=Node(0,None,None,None,"BLACK",0)
>>> root=Node(7,nil,nil,nil,"BLACK",0)
>>> 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",0)
	t.tree_insert(z)

	
>>> t.os_select(t.root,5)
<__main__.Node object at 0x0337DF90>
>>> t.os_select(t.root,5).key
7

>>> t.os_rank(t.tree_search(t.root,6))
4