33import random
44from typing import Any
55
6+
67class MyQueue :
78 def __init__ (self ) -> None :
89 self .data : list [Any ] = []
@@ -18,57 +19,69 @@ def push(self, data: Any) -> None:
1819 def pop (self ) -> Any :
1920 return self .data [self .head ]
2021
22+
2123class MyNode :
2224 __slots__ = ("data" , "left" , "right" , "height" )
23-
25+
2426 def __init__ (self , data : Any ) -> None :
2527 self .data = data
2628 self .left : MyNode | None = None
2729 self .right : MyNode | None = None
2830 self .height = 1
2931
30- def get_height (node : MyNode | None ) -> int :
32+
33+ def get_height (node : MyNode | None ) -> int :
3134 return node .height if node else 0
3235
33- def my_max (a : int , b : int ) -> int :
36+
37+ def my_max (a : int , b : int ) -> int :
3438 return a if a > b else b
3539
40+
3641def right_rotation (node : MyNode ) -> MyNode :
3742 left_child = node .left
3843 if left_child is None :
3944 return node
40-
45+
4146 node .left = left_child .right
4247 left_child .right = node
4348 node .height = my_max (get_height (node .right ), get_height (node .left )) + 1
44- left_child .height = my_max (get_height (left_child .right ), get_height (left_child .left )) + 1
49+ left_child .height = (
50+ my_max (get_height (left_child .right ), get_height (left_child .left )) + 1
51+ )
4552 return left_child
4653
54+
4755def left_rotation (node : MyNode ) -> MyNode :
4856 right_child = node .right
4957 if right_child is None :
5058 return node
51-
59+
5260 node .right = right_child .left
5361 right_child .left = node
5462 node .height = my_max (get_height (node .right ), get_height (node .left )) + 1
55- right_child .height = my_max (get_height (right_child .right ), get_height (right_child .left )) + 1
63+ right_child .height = (
64+ my_max (get_height (right_child .right ), get_height (right_child .left )) + 1
65+ )
5666 return right_child
5767
68+
5869def lr_rotation (node : MyNode ) -> MyNode :
5970 if node .left :
6071 node .left = left_rotation (node .left )
6172 return right_rotation (node )
6273
74+
6375def rl_rotation (node : MyNode ) -> MyNode :
6476 if node .right :
6577 node .right = right_rotation (node .right )
6678 return left_rotation (node )
6779
80+
6881def insert_node (node : MyNode | None , data : Any ) -> MyNode | None :
6982 if node is None :
7083 return MyNode (data )
71-
84+
7285 if data < node .data :
7386 node .left = insert_node (node .left , data )
7487 if get_height (node .left ) - get_height (node .right ) == 2 :
@@ -83,19 +96,21 @@ def insert_node(node: MyNode | None, data: Any) -> MyNode | None:
8396 node = rl_rotation (node )
8497 else :
8598 node = left_rotation (node )
86-
99+
87100 node .height = my_max (get_height (node .right ), get_height (node .left )) + 1
88101 return node
89102
103+
90104def get_left_most (root : MyNode ) -> Any :
91105 while root .left :
92106 root = root .left
93107 return root .data
94108
109+
95110def del_node (root : MyNode | None , data : Any ) -> MyNode | None :
96111 if root is None :
97112 return None
98-
113+
99114 if data == root .data :
100115 if root .left and root .right :
101116 root .data = get_left_most (root .right )
@@ -106,14 +121,14 @@ def del_node(root: MyNode | None, data: Any) -> MyNode | None:
106121 root .left = del_node (root .left , data )
107122 else :
108123 root .right = del_node (root .right , data )
109-
124+
110125 if root .left is None and root .right is None :
111126 root .height = 1
112127 return root
113-
128+
114129 left_height = get_height (root .left )
115130 right_height = get_height (root .right )
116-
131+
117132 if right_height - left_height == 2 :
118133 right_right = get_height (root .right .right ) if root .right else 0
119134 right_left = get_height (root .right .left ) if root .right else 0
@@ -128,35 +143,37 @@ def del_node(root: MyNode | None, data: Any) -> MyNode | None:
128143 root = right_rotation (root )
129144 else :
130145 root = lr_rotation (root )
131-
146+
132147 root .height = my_max (get_height (root .right ), get_height (root .left )) + 1
133148 return root
149+
150+
134151class AVLtree :
135152 __slots__ = ("root" ,)
136-
137- def __init__ (self ) -> None :
153+
154+ def __init__ (self ) -> None :
138155 self .root : MyNode | None = None
139-
140- def get_height (self ) -> int :
156+
157+ def get_height (self ) -> int :
141158 return get_height (self .root )
142-
159+
143160 def insert (self , data : Any ) -> None :
144161 self .root = insert_node (self .root , data )
145-
162+
146163 def del_node (self , data : Any ) -> None :
147164 self .root = del_node (self .root , data )
148-
165+
149166 def __str__ (self ) -> str :
150167 if self .root is None :
151168 return ""
152-
169+
153170 levels = []
154171 queue = [self .root ]
155-
172+
156173 while queue :
157174 current = []
158175 next_level = []
159-
176+
160177 for node in queue :
161178 if node :
162179 current .append (str (node .data ))
@@ -166,26 +183,28 @@ def __str__(self) -> str:
166183 current .append ("*" )
167184 next_level .append (None )
168185 next_level .append (None )
169-
186+
170187 if any (node is not None for node in next_level ):
171188 levels .append (" " .join (current ))
172189 queue = next_level
173190 else :
174191 break
175-
176- return "\n " .join (levels ) + "\n " + "*" * 36
192+
193+ return "\n " .join (levels ) + "\n " + "*" * 36
194+
177195
178196def test_avl_tree () -> None :
179197 t = AVLtree ()
180198 lst = list (range (10 ))
181199 random .shuffle (lst )
182-
200+
183201 for i in lst :
184202 t .insert (i )
185-
203+
186204 random .shuffle (lst )
187205 for i in lst :
188206 t .del_node (i )
189207
208+
190209if __name__ == "__main__" :
191210 test_avl_tree ()
0 commit comments