Skip to content

Commit 5ea1b4a

Browse files
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
1 parent 1ac8d5b commit 5ea1b4a

1 file changed

Lines changed: 48 additions & 29 deletions

File tree

data_structures/binary_tree/avl_tree.py

Lines changed: 48 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import random
44
from typing import Any
55

6+
67
class 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+
2123
class 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+
3641
def 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+
4755
def 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+
5869
def lr_rotation(node: MyNode) -> MyNode:
5970
if node.left:
6071
node.left = left_rotation(node.left)
6172
return right_rotation(node)
6273

74+
6375
def rl_rotation(node: MyNode) -> MyNode:
6476
if node.right:
6577
node.right = right_rotation(node.right)
6678
return left_rotation(node)
6779

80+
6881
def 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+
90104
def get_left_most(root: MyNode) -> Any:
91105
while root.left:
92106
root = root.left
93107
return root.data
94108

109+
95110
def 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+
134151
class 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

178196
def 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+
190209
if __name__ == "__main__":
191210
test_avl_tree()

0 commit comments

Comments
 (0)