Skip to content

Commit c91d0ac

Browse files
author
Nick
authored
AVL_tree.py
1 parent 38f7123 commit c91d0ac

File tree

1 file changed

+294
-0
lines changed

1 file changed

+294
-0
lines changed

‎.py‎

Lines changed: 294 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,294 @@
1+
import random
2+
3+
4+
class Node:
5+
"""Class to create Nodes in AVL Tree"""
6+
7+
__slots__ = 'element', 'left_tree', 'right_tree', 'height'
8+
9+
def __init__(self, element, left_tree, right_tree):
10+
self.element = element
11+
self.left_tree = left_tree
12+
self.right_tree = right_tree
13+
self.height = 0
14+
15+
16+
class AVLTree:
17+
"""AVL tree class with the following methods:
18+
- create nodes in the tree (the tree creates nodes and checks it self
19+
to keep the tree balanced);
20+
- find height of the tree;
21+
- search nodes in the tree;
22+
- display all nodes in the tree;
23+
- delete nodes in the tree (keeping the tree balanced);
24+
- display number of all nodes in the tree;
25+
"""
26+
27+
def __init__(self):
28+
self.root = None
29+
30+
def create_tree(self, root, e):
31+
if root is None:
32+
n = Node(e, None, None)
33+
root = n
34+
else:
35+
if e == root.element:
36+
print(f"""The element {e} is already exist in the tree!""")
37+
return
38+
elif e < root.element:
39+
root.left_tree = self.create_tree(root.left_tree, e)
40+
else:
41+
root.right_tree = self.create_tree(root.right_tree, e)
42+
root.height = self.balance_factor(root)
43+
# If you want to see balance factor for each node
44+
# print('Balance factor for element:', root.element, 'is', root.height)
45+
if root.height < -1:
46+
balanced_left_tree = self.left_rotation(root)
47+
return balanced_left_tree
48+
if root.height > 1:
49+
balanced_right_tree = self.right_rotation(root)
50+
return balanced_right_tree
51+
return root
52+
53+
def height_of_tree(self, root):
54+
if root is None:
55+
return 0
56+
else:
57+
height_left_tree = 1 + self.height_of_tree(root.left_tree)
58+
height_right_tree = 1 + self.height_of_tree(root.right_tree)
59+
return max(height_left_tree, height_right_tree)
60+
61+
def final_height(self):
62+
if self.root is None:
63+
return 0
64+
height = (self.height_of_tree(self.root)) - 1
65+
print('The height of the tree is', height)
66+
return height
67+
68+
def balance_factor(self, root):
69+
if root is None:
70+
return 0
71+
return self.height_of_tree(root.right_tree) - self.height_of_tree(root.left_tree)
72+
73+
def left_rotation(self, root):
74+
parent = root
75+
child = parent.left_tree
76+
77+
# left-left rotation
78+
if parent == self.root and parent.height == -2 and child.height == -1:
79+
parent.left_tree = child.right_tree
80+
child.right_tree = parent
81+
self.root = child
82+
parent.height = self.balance_factor(child)
83+
child.height = self.balance_factor(child)
84+
return child
85+
86+
elif parent != self.root and parent.height == -2 and child.height == -1:
87+
parent.left_tree = child.right_tree
88+
child.right_tree = parent
89+
parent.height = self.balance_factor(child)
90+
child.height = self.balance_factor(child)
91+
return child
92+
93+
# left-right rotation
94+
elif parent == self.root and parent.height == -2 and child.height == 1:
95+
ptr_ptr = child.right_tree
96+
child.right_tree = ptr_ptr.left_tree
97+
parent.left_tree = ptr_ptr.right_tree
98+
ptr_ptr.left_tree = child
99+
ptr_ptr.right_tree = parent
100+
self.root = ptr_ptr
101+
parent.height = self.balance_factor(child)
102+
child.height = self.balance_factor(child)
103+
ptr_ptr.height = self.balance_factor(ptr_ptr)
104+
return ptr_ptr
105+
106+
elif parent != self.root and parent.height == -2 and child.height == 1:
107+
ptr_ptr = child.right_tree
108+
child.right_tree = ptr_ptr.left_tree
109+
parent.left_tree = ptr_ptr.right_tree
110+
ptr_ptr.left_tree = child
111+
ptr_ptr.right_tree = parent
112+
parent.height = self.balance_factor(child)
113+
child.height = self.balance_factor(child)
114+
ptr_ptr.height = self.balance_factor(ptr_ptr)
115+
return ptr_ptr
116+
117+
# case for deletion function
118+
elif parent == self.root and parent.height == -2 and child.height == 0:
119+
parent.left_tree = child.right_tree
120+
child.right_tree = parent
121+
self.root = child
122+
parent.height = self.balance_factor(parent)
123+
child.height = self.balance_factor(child)
124+
return child
125+
126+
elif parent != self.root and parent.height == -2 and child.height == 0:
127+
parent.left_tree = child.right_tree
128+
child.right_tree = parent
129+
parent.height = self.balance_factor(parent)
130+
child.height = self.balance_factor(child)
131+
return child
132+
133+
def right_rotation(self, root):
134+
parent = root
135+
child = parent.right_tree
136+
137+
# right-right rotation
138+
if parent == self.root and parent.height == 2 and child.height == 1:
139+
parent.right_tree = child.left_tree
140+
child.left_tree = parent
141+
self.root = child
142+
parent.height = self.balance_factor(child)
143+
child.height = self.balance_factor(child)
144+
return child
145+
146+
elif parent != self.root and parent.height == 2 and child.height == 1:
147+
parent.right_tree = child.left_tree
148+
child.left_tree = parent
149+
parent.height = self.balance_factor(child)
150+
child.height = self.balance_factor(child)
151+
return child
152+
153+
# right-left rotation
154+
elif parent == self.root and parent.height == 2 and child.height == -1:
155+
ptr_ptr = child.left_tree
156+
child.left_tree = ptr_ptr.right_tree
157+
parent.right_tree = ptr_ptr.left_tree
158+
ptr_ptr.left_tree = parent
159+
ptr_ptr.right_tree = child
160+
self.root = ptr_ptr
161+
parent.height = self.balance_factor(child)
162+
child.height = self.balance_factor(child)
163+
ptr_ptr.height = self.balance_factor(ptr_ptr)
164+
return ptr_ptr
165+
166+
elif parent != self.root and parent.height == 2 and child.height == -1:
167+
ptr_ptr = child.left_tree
168+
child.left_tree = ptr_ptr.right_tree
169+
parent.right_tree = ptr_ptr.left_tree
170+
ptr_ptr.left_tree = parent
171+
ptr_ptr.right_tree = child
172+
parent.height = self.balance_factor(child)
173+
child.height = self.balance_factor(child)
174+
ptr_ptr.height = self.balance_factor(ptr_ptr)
175+
return ptr_ptr
176+
177+
# case for deletion function
178+
elif parent == self.root and parent.height == 2 and child.height == 0:
179+
parent.right_tree = child.left_tree
180+
child.left_tree = parent
181+
self.root = child
182+
parent.height = self.balance_factor(parent)
183+
child.height = self.balance_factor(child)
184+
return child
185+
186+
elif parent != self.root and parent.height == 2 and child.height == 0:
187+
parent.right_tree = child.left_tree
188+
child.left_tree = parent
189+
parent.height = self.balance_factor(parent)
190+
child.height = self.balance_factor(child)
191+
return child
192+
193+
def find_greatest(self, root):
194+
if root is None:
195+
return None
196+
else:
197+
if root.right_tree is not None:
198+
return self.find_greatest(root.right_tree)
199+
return root
200+
201+
def search_element(self, root, e):
202+
if root is not None:
203+
if e == root.element:
204+
print()
205+
print(f"""Element {e} is found in the tree!""")
206+
return True
207+
elif e < root.element:
208+
self.search_element(root.left_tree, e)
209+
else:
210+
self.search_element(root.right_tree, e)
211+
else:
212+
print()
213+
print(f"""Element {e} is {'not found'.upper()} in the tree""")
214+
return False
215+
216+
def removing_element(self, root, e):
217+
if root is None:
218+
print(f"""Element {e} is {'not'.upper()} found in the tree""")
219+
return root
220+
else:
221+
if root.left_tree is not None:
222+
root.left_tree = self.removing_element(root.left_tree, e)
223+
if root.element == e:
224+
p = root
225+
if root.left_tree is None and root.right_tree is None:
226+
p = None
227+
root = p
228+
print(f"""Element {e} was deleted from the tree""")
229+
return root
230+
elif root.left_tree is not None and root.right_tree is None:
231+
root = root.left_tree
232+
p = root
233+
print(f"""Element {e} was deleted from the tree""")
234+
return p
235+
elif root.left_tree is None and root.right_tree is not None:
236+
root = root.right_tree
237+
p = root
238+
print(f"""Element {e} was deleted from the tree""")
239+
return p
240+
elif root.left_tree is not None and root.right_tree is not None:
241+
p = self.find_greatest(root.left_tree)
242+
root.element, p.element = p.element, root.element
243+
root.left_tree = self.removing_element(root.left_tree, e)
244+
if root.right_tree is not None:
245+
root.right_tree = self.removing_element(root.right_tree, e)
246+
root.height = self.balance_factor(root)
247+
if root.height < -1:
248+
return self.left_rotation(root)
249+
elif root.height > 1:
250+
return self.right_rotation(root)
251+
return root
252+
253+
def display_in_order(self, root):
254+
if root is not None:
255+
self.display_in_order(root.left_tree)
256+
print(root.element, end=' ')
257+
self.display_in_order(root.right_tree)
258+
259+
def count_number_of_nodes(self, root):
260+
if root is not None:
261+
q = self.count_number_of_nodes(root.left_tree)
262+
w = self.count_number_of_nodes(root.right_tree)
263+
return q + w + 1
264+
return 0
265+
266+
267+
testing_array = [0] * 5000
268+
269+
for i in range(len(testing_array)):
270+
testing_array[i] = random.randint(-1000000, 1000000)
271+
272+
# here again I use set, to create unique values in AVL tree since AVL tree have to have unique keys
273+
unique_values = list(set(testing_array))
274+
275+
x = AVLTree()
276+
x.root = x.create_tree(x.root, unique_values[0])
277+
278+
for j in range(1, len(unique_values)):
279+
x.create_tree(x.root, unique_values[j])
280+
281+
print(unique_values)
282+
283+
x.display_in_order(x.root)
284+
print()
285+
x.final_height()
286+
print(x.balance_factor(x.root))
287+
288+
for k in range(len(unique_values)):
289+
x.removing_element(x.root, unique_values[k])
290+
# Uncomment this if you want to see height, balance factor for each node and
291+
# number of nodes in the tree after each deletion operation
292+
# x.final_height()
293+
# print('Balance factor:', x.balance_factor(x.root))
294+
# print('Number of nodes in the tree:', x.count_number_of_nodes(x.root))

0 commit comments

Comments
 (0)