-
Notifications
You must be signed in to change notification settings - Fork 0
/
red_black_tree.py
168 lines (146 loc) · 4.66 KB
/
red_black_tree.py
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
import sys
BLACK = 0
RED = 1
class Node:
def __init__(self, data):
self.data = data
self.color = RED
self.parent = None
self.right = None
self.left = None
class RedBlackTree:
def __init__(self):
self.TNULL = Node(0)
self.TNULL.color = BLACK
self.TNULL.left = None
self.TNULL.right = None
self.root = self.TNULL
# rotate left at node x
def left_rotate(self, x):
y = x.right
x.right = y.left
if y.left != self.TNULL:
y.left.parent = x
y.parent = x.parent
if not x.parent:
self.root = y
elif x == x.parent.left:
x.parent.left = y
else:
x.parent.right = y
y.left = x
x.parent = y
# rotate right at node x
def right_rotate(self, x):
y = x.left
x.left = y.right
if y.right != self.TNULL:
y.right.parent = x
y.parent = x.parent
if not x.parent:
self.root = y
elif x == x.parent.right:
x.parent.right = y
else:
x.parent.left = y
y.right = x
x.parent = y
# insert the key to the tree in its appropriate position
# and fix the tree
def insert(self, data):
# Ordinary Binary Search Insertion
node = Node(data)
node.parent = None
node.left = self.TNULL
node.right = self.TNULL
node.color = RED
y = None
x = self.root
while x != self.TNULL:
y = x
if node.data < x.data:
x = x.left
else:
x = x.right
# y is parent of x
node.parent = y
if not self.root:
self.root = node
elif node.data < y.data:
y.left = node
else:
y.right = node
# if new node is a root node, simply return
if not node.parent:
node.color = BLACK
return
# if the grandparent is None, simply return
if not node.parent.parent:
return
# Fix the tree
self.__fix_insert(node)
# fix the red-black tree
def __fix_insert(self, node):
while node.parent.color == RED:
# check if parent is right child
if node.parent == node.parent.parent.right:
uncle = node.parent.parent.left
if uncle.color == RED:
# case 3.1
uncle.color = BLACK
node.parent.color = BLACK
node.parent.parent.color = RED
node = node.parent.parent
else:
if node == node.parent.left:
# case 3.2.2 (triangle)
node = node.parent
self.right_rotate(node)
# case 3.2.1 (line)
node.parent.color = BLACK
node.parent.parent.color = RED
self.left_rotate(node.parent.parent)
else:
uncle = node.parent.parent.right
if uncle.color == RED:
# mirror case 3.1
uncle.color = BLACK
node.parent.color = BLACK
node.parent.parent.color = RED
node = node.parent.parent
else:
if node == node.parent.right:
# mirror case 3.2.2 (triangle)
node = node.parent
self.left_rotate(node)
# mirror case 3.2.1(line)
node.parent.color = BLACK
node.parent.parent.color = RED
self.right_rotate(node.parent.parent)
if node == self.root:
break
self.root.color = BLACK
def __print_helper(self, node, indent, last):
# print the tree structure on the screen
if node != self.TNULL:
sys.stdout.write(indent)
if last:
sys.stdout.write("R----")
indent += " "
else:
sys.stdout.write("L----")
indent += "| "
s_color = "RED" if node.color == 1 else "BLACK"
print(str(node.data) + "(" + s_color + ")")
self.__print_helper(node.left, indent, False)
self.__print_helper(node.right, indent, True)
# print the tree structure on the screen
def pretty_print(self):
self.__print_helper(self.root, "", True)
tree = RedBlackTree()
tree.insert(1)
tree.insert(2)
tree.insert(6)
tree.insert(4)
tree.insert(5)
tree.pretty_print()