Using Python to Write a Printable Red-black Tree

Posted on

引文

红黑树是计算机算法中非常经典的一种数据结构。

然而之前由于个人学习的疏忽，一直无缘了解什么是红黑树（当时学习了AVL树）。

如今，回过头来学习红黑树，深深感受到这个算法的美妙，并震惊于发明者的聪明才智。

已经有很多人就红黑树的算法进行了解释，此处不做过多赘述。这里，我简单的附上一个可以打印的红黑树的代码源码，供读者参考：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
class TreeNode(object):
    def __init__(self, key, color=None, left=None, right=None, parent=None):
        self.key = key
        self.color = color # 0 for red, 1 for black
        self.left = left
        self.right = right
        self.parent = parent

    def __str__(self, depth=0):
        ret = ""
        # Print right branch
        if self.right.key !=  -1:
            ret += self.right.__str__(depth + 1)
        # Print own value
        if self.color == 0:
            pcolor = 'r'
        else:
            pcolor = 'b'
        ret += "\n" + ("    "*depth) + str(self.key) + pcolor
        # Print left branch
        if self.left.key !=  -1:
            ret += self.left.__str__(depth + 1)

        return ret

class Tree(object):
    def __init__(self):
        self.nil = TreeNode(-1, 1)
        self.root = self.nil

def leftRotate(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(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(T, z):
    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
    else:
        if z.key < y.key:
            y.left = z
        else:
            y.right = z
    z.left = z.right = T.nil
    z.color = 0
    rbInsertFix(T, z)

def rbInsertFix(T, z):
    while z.parent.color == 0:
        if z.parent == z.parent.parent.left:
            y = z.parent.parent.right
            if y.color == 0:
                z.parent.color = 1
                y.color = 1
                z.parent.parent.color = 0
                z = z.parent.parent
            else:
                if z == z.parent.right:
                    z = z.parent
                    leftRotate(T, z)
                z.parent.color = 1
                z.parent.parent.color = 0
                rightRotate(T, z.parent.parent)
        else:
            y = z.parent.parent.left
            if y.color == 0:
                z.parent.color = 1
                y.color = 1
                z.parent.parent.color = 0
                z = z.parent.parent
            else:
                if z == z.parent.left:
                    z = z.parent
                    rightRotate(T, z)
                z.parent.color = 1
                z.parent.parent.color = 0
                leftRotate(T, z.parent.parent)
            
    T.root.color = 1


def treeMin(T, x):
    while x.left != T.nil:
        x = x.left
    return x


def treeMax(T, x):
    while x.right != T.nil:
        x = x.right
    return x


def treeSuccessor(T, x):
    if x.right != T.nil:
        return treeMin(T, x.right)
    y = x.parent
    while y != T.nil and x == y.right:
        x = y
        y = x.parent
    return y


def treePreSuccessor(T, x):
    if x.left != T.nil:
        return treeMax(T, x.left)
    y = x.parent
    while y != T.nil and x == y.left:
        x = y
        y = x.parent
    return y

def rbDelete(T, z):
    if z.left == T.nil or z.right == T.nil:
        y = z
    else:
        y = treeSuccessor(T, z)
    if y.left != T.nil:
        x = y.left
    else:
        x = y.right
    x.parent = y.parent
    if y.parent == T.nil:
        T.root = x
    else:
        if y == y.parent.left:
            y.parent.left = x
        else:
            y.parent.right = x
    if y != z:
        z.key = y.key
    if y.color == 1:
        rbDeleteFix(T, x)
    return y

def rbDeleteFix(T, x):
    while x != T.root and x.color == 1:
        if x == x.parent.left:
            w = x.parent.right
            if w.color == 0:
                w.color = 1
                x.parent.color = 0
                leftRotate(T, x.parent)
                w = x.parent.right
            if w.left.color == 1 and w.right.color == 1:
                w.color = 0
                x = x.parent
            else:
                if w.right.color == 1:
                    w.left.color = 1
                    w.color = 0
                    rightRotate(T, w)
                    w = x.parent.right
                w.color = x.parent.color
                x.parent.color = 1
                w.right.color = 1
                leftRotate(T, x.parent)
                x = T.root
        else:
            w = x.parent.left
            if w.color == 0:
                w.color = 1
                x.parent.color = 0
                rightRotate(T, x.parent)
                w = x.parent.left
            if w.right.color == 1 and w.left.color == 1:
                w.color = 0
                x = x.parent
            else:
                if w.left.color == 1:
                    w.right.color = 1
                    w.color = 0
                    leftRotate(T, w)
                    w = x.parent.left
                w.color = x.parent.color
                x.parent.color = 1
                w.left.color = 1
                rightRotate(T, x.parent)
                x = T.root
    x.color = 1

def main():
    T = Tree()
    rbInsert(T, TreeNode(11))
    rbInsert(T, TreeNode(2))
    rbInsert(T, TreeNode(14))
    rbInsert(T, TreeNode(1))
    # rbDelete(T, T.root)
    rbInsert(T, TreeNode(7))
    rbInsert(T, TreeNode(15))
    rbInsert(T, TreeNode(5))
    rbInsert(T, TreeNode(8))
    rbInsert(T, TreeNode(4))
    # rbInsert(T, TreeNode(7))
    rbDelete(T, T.root.right.left)
    print(T.root)


if __name__ == '__main__':
    main()

'''
Output Example:
            15r
        14b
    11r
        8b
7b
        5b
            4r
    2r
        1b
'''