1. 程式人生 > >[從頭學數學] 第261節 Python實現資料結構:紅黑樹(RB Tree)

[從頭學數學] 第261節 Python實現資料結構:紅黑樹(RB Tree)

劇情提要:
阿偉看到了一本比較有趣的書,是關於《計算幾何》的,2008年由北清派出版。很好奇
它裡面講了些什麼,就來看看啦。


正劇開始:
星曆2016年09月09日 11:26:00, 銀河系厄爾斯星球中華帝國江南行省。

[工程師阿偉]正在和[機器小偉]一起研究[計算幾何]]。



<span style="font-size:18px;">#
from random import randint

RED = 'R'
BLACK = 'B'

###
# @usage   紅黑樹
# @author  mw
# @date    2016年09月09日  星期五  11:24:48 
# @param
# @return
#
###
class RBTree():        
    class __RBTreeNode():
        '''紅黑樹的節點型別'''
        def __init__(self, val):
            self.val = val
            self.left = None
            self.right = None
            self.parent = None
            self.color = '*';

        def paint(self, color):
            self.color = color
            
        def __iter__(self):
            if self.left != None:
                for elem in self.left:
                    yield elem

            if (self.val != None):
                yield self.val

            if self.right != None:
                for elem in self.right:
                    yield elem

        #迭代的是Node型別,用於刪除結點
        def iternodes(self):
            if self.left != None:
                for elem in self.left.iternodes():
                    yield elem

            if self != None and self.val != None:
                yield self

            if self.right != None:
                for elem in self.right.iternodes():
                    yield elem

        def info(self):
            s = 'Key='+str(self.val)+', '+\
                'LChild='+str(self.left)+', '+\
                'RChild='+str(self.right)+', '+\
                'parent='+str(self.parent)+', '+\
                'Color='+str(self.color);

            print(s);
            
        def __str__(self):
            return str(self.val);

        def __repr__(self):
            if self != None:
                s_1 = str(self.val);
            else:
                s_1 = 'None';

            if self.left != None:
                s_2 = str(self.left.val);
            else:
                s_2 = 'None';

            if self.right != None:
                s_3 = str(self.right.val);
            else:
                s_3 = 'None';

            s_4 = str(self.parent);
            s_5 = str(self.color);
                
            return '__RBTreeNode('+s_1+', ' + s_2 +', ' + s_3 +', '+ s_4 +', '+s_5+')';
            
            
    def __init__(self):
       # self.items = []
        self.root = None
        self.zlist = []

    def LRotate(self, x):
        # x是一個RBTree.__RBTreeNode
        y = x.right
        if y is None:
            # 右節點為空,不旋轉
            pass;
        else:
            beta = y.left
            x.right = beta
            if beta is not None:
                beta.parent = x

            p = x.parent
            y.parent = p
            if p is None:
                # x原來是root
                self.root = y
            elif x == p.left:
                p.left = y
            else:
                p.right = y
            y.left = x
            x.parent = y
 

    def RRotate(self, y):
        # y是一個節點
        x = y.left
        if x is None:
            # 左節點為空,不旋轉
            pass;
        else:
            beta = x.right
            y.left = beta
            if beta is not None:
                beta.parent = y

            p = y.parent
            x.parent = p
            if p is None:
                # y原來是root
                self.root = x
            elif y == p.left:
                p.left = x
            else:
                p.right = x
            x.right = y
            y.parent = x


    #插入
    def insert(self, val):
        z = RBTree.__RBTreeNode(val)
        y = None
        x = self.root
        while x is not None:
            y = x
            if z.val < x.val:
                x = x.left
            else:
                x = x.right


        z.paint(RED)
        z.parent = y

        if y is None:
            # 插入z之前為空的RBTree
            self.root = z
            self.insert_Fixup(z)
            return

        if z.val < y.val:
            y.left = z
        else:
            y.right = z

        if y.color == RED:
            # z的父節點y為紅色,需要fixup。
            # 如果z的父節點y為黑色,則不用調整
            self.insert_Fixup(z)

        else:
            return

    def insert_Fixup(self, z):
        # case 1:z為root節點
        if z.parent is None:
            z.paint(BLACK)
            self.root = z
            return

        # case 2:z的父節點為黑色
        if z.parent.color == BLACK:
            # 包括了z處於第二層的情況
            # 這裡感覺不必要啊。。似乎z.parent為黑色則不會進入fixup階段
            return

        # 下面的幾種情況,都是z.parent.color == RED:
        # 節點y為z的uncle
        p = z.parent
        g = p.parent  # g為x的grandpa
        if g is None:
            return
            #   return 這裡不能return的。。。
        if g.right == p:
            y = g.left
        else:
            y = g.right

        # case 3-0:z沒有叔叔。即:y為NIL節點
        # 注意,此時z的父節點一定是RED
        if y == None:
            if z == p.right and p == p.parent.left:
                # 3-0-0:z為右兒子,且p為左兒子,則把p左旋
                # 轉化為3-0-1或3-0-2的情況
                self.LRotate(p)
                p, z = z, p
                g = p.parent
            elif z == p.left and p == p.parent.right:
                self.RRotate(p)
                p, z = z, p

            g.paint(RED)
            p.paint(BLACK)
            if p == g.left:
                # 3-0-1:p為g的左兒子
                self.RRotate(g)
            else:
                # 3-0-2:p為g的右兒子
                self.LRotate(g)

            return

        # case 3-1:z有黑叔
        elif y.color == BLACK:
            if p.right == z and p.parent.left == p:
                # 3-1-0:z為右兒子,且p為左兒子,則左旋p
                # 轉化為3-1-1或3-1-2
                self.LRotate(p)
                p, z = z, p
            elif p.left == z and p.parent.right == p:
                self.RRotate(p)
                p, z = z, p

            p = z.parent
            g = p.parent

            p.paint(BLACK)
            g.paint(RED)
            if p == g.left:
                # 3-1-1:p為g的左兒子,則右旋g
                self.RRotate(g)
            else:
                # 3-1-2:p為g的右兒子,則左旋g
                self.LRotate(g)

            return

        # case 3-2:z有紅叔
        # 則塗黑父和叔,塗紅爺,g作為新的z,遞迴呼叫
        else:
            y.paint(BLACK)
            p.paint(BLACK)
            g.paint(RED)
            new_z = g
            self.insert_Fixup(new_z)
            
    #刪除
    def delete(self, val):
        curNode = self.root
        while curNode is not None:
            if val < curNode.val:
                curNode = curNode.left
            elif val > curNode.val:
                curNode = curNode.right
            else:
                # 找到了值為val的元素,正式開始刪除
                if curNode.left is None and curNode.right is None:
                    # case1:curNode為葉子節點:直接刪除即可
                    if curNode == self.root:
                        self.root = None
                    else:
                        p = curNode.parent
                        if curNode == p.left:
                            p.left = None
                        else:
                            p.right = None

                elif curNode.left is not None and curNode.right is not None:
                    sucNode = self.SUCCESOR(curNode)
                    curNode.val, sucNode.val  = sucNode.val, curNode.val
                    self.delete(sucNode.val)

                else:
                    p = curNode.parent
                    if curNode.left is None:
                        x = curNode.right
                    else:
                        x = curNode.left
                    if curNode == p.left:
                        p.left = x
                    else:
                        p.right = x
                    x.parent = p
                    if curNode.color == BLACK:
                        self.delete_Fixup(x)
                curNode = None
        return False

    def delete_Fixup(self, x):
        p = x.parent
        # w:x的兄弟結點
        if x == p.left:
            w = x.right
        else:
            w = x.left

        # case1:x的兄弟w是紅色的
        if w.color == RED:
            p.paint(RED)
            w.paint(BLACK)
            if w == p.right:
                self.LRotate(p)
            else:
                self.RRotate(p)

        if w.color == BLACK:
            # case2:x的兄弟w是黑色的,而且w的兩個孩子都是黑色的
            if w.left.color == BLACK and w.right.color == BLACK:
                w.paint(RED)
                if p.color == BLACK:
                    return
                else:
                    p.color = BLACK
                    self.delete_Fixup(p)

            # case3:x的兄弟w是黑色的,而且w的左兒子是紅色的,右兒子是黑色的
            if w.left.color == RED and w.color == BLACK:
                w.left.paint(BLACK)
                w.paint(RED)
                self.RRotate(w)

            # case4:x的兄弟w是黑色的,而且w的右兒子是紅
            if w.right.color == RED:
                p.paint(BLACK)
                w.paint(RED)
                if w == p.right:
                    self.LRotate(p)
                else:
                    self.RRotate(p)

    #紅黑樹資訊查詢
    def info(self):
        a = [];
        for x in self:
            a.append(x);

        print(a);

    def __iter__(self):
        if self.root != None:
            return self.root.__iter__()
        else:
            return [].__iter__()

    #傳回結點的原始資訊
    def iternodes(self):
        if self.root != None:
            return self.root.iternodes()
        else:
            return [None];
            

    #查詢,返回的是結點
    def find(self, key):
        def _find(key,  node):
            if node is None:
                return None
            elif key < node.val:
                return _find(key, node.left)
            elif key > node.val:
                return _find(key, node.right)
            else:
                return node
            
        if self.root is None:
            return None
        else:
            return _find(key, self.root)

    #找最小元素
    def findMin(self):
        def _findMin(node):
            if node.left:
                return _findMin(node.left)
            else:
                return node
            
        if self.root is None:
            return None
        else:
            return _findMin(self.root)        
    

    
    #找最大元素
    def findMax(self):
        def _findMax(node):
            if node.right:
                return _findMax(node.right)
            else:
                return node
        
        if self.root is None:
            return None
        else:
            return _findMax(self.root)

    #求結點高度
    def height(self, node):
        if (node == None):
            return 0;
        else:
            m = self.height(node.left);
            n = self.height(node.right);
            return max(m, n)+1;

    #尋找節點路徑
    def findNodePath(self, root, node):
        path = [];
        if root == None or root.val == None:
            path = [];
            return path
            
        while (root != node):
            if node.val < root.val:
                path.append(root);
                root = root.left;
            elif node.val >= root.val:
                path.append(root);
                root = root.right;
            else:
                break;

        path.append(root);
        return path;

    #尋找父結點
    def parent(self, root, node):
        path = self.findNodePath(root, node);
        if (len(path)>1):
            return path[-2];
        else:
            return None;

    #求某元素是在樹的第幾層
    #約定根為0層
    #這個計算和求結點的Height是不一樣的
    def level(self, elem):
        if self.root != None:
            node = self.root;
            lev = 0;

            while (node != None):
                if elem < node.val:
                    node = node.left;
                    lev+=1;
                elif elem > node.val:
                    node = node.right;
                    lev+=1;
                else:
                    return lev;

            return -1;

        else:
            return -1;

    def __len__(self):
        a = [];
        for x in self:
            a.append(x);

        return len(a);
    
    def __contains__(self, val):
        if self.find(val) != None:
            return True;

        return False;
#</span>

用例:
<span style="font-size:18px;">#
if __name__ == '__main__':
    rbt=RBTree()
    b = []
    N = 10;

    data = [0, 24, 26, 79, 93, 99, 335, 378, 440, 494];

    data2 = [90, 115, 168, 236, 270, 271, 300, 393, 465, 489];
    
    #生成隨機原始資料
    for i in range(N):
        m = randint(0, 500)
        rbt.insert(m)
        b.append(m)
        
    #插入
    for x in data2:
        rbt.insert(x);
    
    for node in rbt.iternodes():
        node.info();

    print('root:');
    rbt.root.info();   

    #刪除
    rbt.delete(236);
    rbt.delete(233);
    print('root:');
    rbt.root.info();
    print(len(rbt));

    #查詢
    p = rbt.find(168);
    if (p != None):
        p.info();

    #
    print('max:');
    rbt.findMax().info();
    print('min:');
    rbt.findMin().info();
    print('height:', rbt.height(rbt.root));
    print(168 in rbt);
    print(1002 in rbt);

    #
    for x in rbt.iternodes():
        print(rbt.findNodePath(rbt.root, x));
        print(rbt.level(x.val));
		
>>> 
Key=33, LChild=None, RChild=None, parent=90, Color=R
Key=90, LChild=33, RChild=115, parent=141, Color=B
Key=115, LChild=None, RChild=None, parent=90, Color=R
Key=141, LChild=90, RChild=191, parent=205, Color=R
Key=168, LChild=None, RChild=None, parent=191, Color=R
Key=191, LChild=168, RChild=None, parent=141, Color=B
Key=205, LChild=141, RChild=256, parent=275, Color=B
Key=236, LChild=None, RChild=None, parent=256, Color=B
Key=256, LChild=236, RChild=270, parent=205, Color=R
Key=270, LChild=None, RChild=271, parent=256, Color=B
Key=271, LChild=None, RChild=None, parent=270, Color=R
Key=275, LChild=205, RChild=456, parent=None, Color=B
Key=300, LChild=None, RChild=None, parent=393, Color=R
Key=393, LChild=300, RChild=416, parent=456, Color=B
Key=416, LChild=None, RChild=None, parent=393, Color=R
Key=456, LChild=393, RChild=465, parent=275, Color=B
Key=461, LChild=None, RChild=None, parent=465, Color=B
Key=465, LChild=461, RChild=495, parent=456, Color=R
Key=489, LChild=None, RChild=None, parent=495, Color=R
Key=495, LChild=489, RChild=None, parent=465, Color=B
root:
Key=275, LChild=205, RChild=456, parent=None, Color=B
root:
Key=275, LChild=205, RChild=456, parent=None, Color=B
19
Key=168, LChild=None, RChild=None, parent=191, Color=R
max:
Key=495, LChild=489, RChild=None, parent=465, Color=B
min:
Key=33, LChild=None, RChild=None, parent=90, Color=R
height: 5
True
False
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(141, 90, 191, 205, R), __RBTreeNode(90, 33, 115, 141, B), __RBTreeNode(33, None, None, 90, R)]
4
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(141, 90, 191, 205, R), __RBTreeNode(90, 33, 115, 141, B)]
3
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(141, 90, 191, 205, R), __RBTreeNode(90, 33, 115, 141, B), __RBTreeNode(115, None, None, 90, R)]
4
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(141, 90, 191, 205, R)]
2
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(141, 90, 191, 205, R), __RBTreeNode(191, 168, None, 141, B), __RBTreeNode(168, None, None, 191, R)]
4
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(141, 90, 191, 205, R), __RBTreeNode(191, 168, None, 141, B)]
3
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B)]
1
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(256, None, 270, 205, R)]
2
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(256, None, 270, 205, R), __RBTreeNode(270, None, 271, 256, B)]
3
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(205, 141, 256, 275, B), __RBTreeNode(256, None, 270, 205, R), __RBTreeNode(270, None, 271, 256, B), __RBTreeNode(271, None, None, 270, R)]
4
[__RBTreeNode(275, 205, 456, None, B)]
0
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(456, 393, 465, 275, B), __RBTreeNode(393, 300, 416, 456, B), __RBTreeNode(300, None, None, 393, R)]
3
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(456, 393, 465, 275, B), __RBTreeNode(393, 300, 416, 456, B)]
2
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(456, 393, 465, 275, B), __RBTreeNode(393, 300, 416, 456, B), __RBTreeNode(416, None, None, 393, R)]
3
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(456, 393, 465, 275, B)]
1
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(456, 393, 465, 275, B), __RBTreeNode(465, 461, 495, 456, R), __RBTreeNode(461, None, None, 465, B)]
3
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(456, 393, 465, 275, B), __RBTreeNode(465, 461, 495, 456, R)]
2
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(456, 393, 465, 275, B), __RBTreeNode(465, 461, 495, 456, R), __RBTreeNode(495, 489, None, 465, B), __RBTreeNode(489, None, None, 495, R)]
4
[__RBTreeNode(275, 205, 456, None, B), __RBTreeNode(456, 393, 465, 275, B), __RBTreeNode(465, 461, 495, 456, R), __RBTreeNode(495, 489, None, 465, B)]
3
>>> 

#</span>

參見http://blog.csdn.net/napo1987/article/details/37928985
參見http://www.cnblogs.com/zjutzz/p/3281319.html

本節到此結束,欲知後事如何,請看下回分解。