-
Notifications
You must be signed in to change notification settings - Fork 0
/
AVL_tree.py
149 lines (123 loc) · 3.54 KB
/
AVL_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
### AVL tree implementation ###
### check https://habrahabr.ru/post/150732/ ###
class Tree:
"""
Interface for working with avl trees
"""
def __init__(self, data):
self.root = Node(data)
def insert(self, data):
if self.root is None:
self.root = Node(data)
else:
self.root = self.root.insert(data)
def print_out(self):
self.root.print_node()
def remove(self, data):
self.root = self.root.remove(data)
class Node(object):
"""
Node object interface
"""
def __init__(self, data):
self.data = data
self.left = None
self.right = None
self.height = 1
def bfactor(self):
"""
:return: difference between right and left subtrees
"""
if self.left is None:
if self.right is None:
return 0
else:
return self.right.height
else:
if self.right is None:
return self.left.height
else:
return self.right.height - self.left.height
def fix_height(self):
"""
Fix height value in current node
"""
if self.left is None:
if self.right is None:
self.height = 1
else:
self.height = self.right.height + 1
else:
if self.right is None:
self.height = self.left.height + 1
else:
self.height = max(self.left.height, self.right.height) + 1
def rotate_right(self):
q = self.left
self.left = self.left.right
q.right = self
self.fix_height()
q.fixheight()
return q
def rotate_left(self):
p = self.right
self.right = self.right.left
p.left = self
self.fix_height()
p.fix_height()
return p
def balance(self):
"""
Function to balance nodes
"""
self.fix_height()
if self.bfactor() == 2:
if self.right is None:
return self.rotate_right()
if self.right.bfactor()<0:
self.right = self.right.rotateright()
return self.rotate_left()
if self.bfactor() == -2:
if self.left is None:
return self.rotate_left()
if self.left.bfactor() > 0:
self.left = self.left.rotateleft()
return self.rotate_right()
return self
def insert(self, data):
"""
Insertion a node with data
"""
if data >= self.data:
if self.right is None:
self.right = Node(data)
self.height = 2
return self
else:
self.right = self.right.insert(data)
if data < self.data:
if self.left is None:
self.left = Node(data)
self.height = 2
return self
else:
self.left = self.left.insert(data)
return self.balance()
def print_node(self, level=0):
if self.right:
self.right.print_node(level + 1)
print('\t' * level, ' /')
print('\t' * level, self.data)
if self.left:
print('\t' * level, ' \\')
self.left.print_node(level + 1)
if __name__ == "__main__":
tree = Tree(10)
tree.insert(20)
tree.insert(5)
tree.insert(25)
tree.insert(30)
tree.insert(1)
tree.insert(45)
tree.insert(60)
tree.print_out()