-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathavl.py
152 lines (125 loc) · 5.18 KB
/
avl.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
from typing import Optional, Type
from generic import ABinarySearchTree, BinaryNode
class AvlNode(BinaryNode):
""" Object extends BinaryNode specific for AvlTree """
nxt: 'AvlNode'
prev: 'AvlNode'
def update(self):
self.depth = max(self.left.depth, self.right.depth) + 1
@property
def balance(self):
""" Return balance of tree
-1 -- right subtree has bigger depth
0 -- left and right subtree has same depth
1 -- left subtree has bigger depth
"""
return self.left.depth - self.right.depth
class AvlTree(ABinarySearchTree):
EXTERNAL_NODE = AvlNode(0, depth=0, external=True)
EXTERNAL_NODE.left = EXTERNAL_NODE
EXTERNAL_NODE.right = EXTERNAL_NODE
def __init__(self):
self.root = self.EXTERNAL_NODE
def insert(self, key, value) -> None:
self.root = self._insert(key, value, self.root)
def _insert(self, key, value, node: AvlNode) -> AvlNode:
if node.external:
new_node = AvlNode(key, value, left=self.EXTERNAL_NODE, right=self.EXTERNAL_NODE)
self._add_node_to_chain(new_node, self._find_lower(self.root, key),
self._find_bigger(self.root, key))
return new_node
if key == node.key:
# key is alread in tree - nothing to do
return node
elif key < node.key:
node.left = self._insert(key, value, node.left)
else:
node.right = self._insert(key, value, node.right)
left_depth = node.left.depth
right_depth = node.right.depth
if abs(left_depth - right_depth) <= 1:
# everything is OK
pass
else:
if left_depth > right_depth:
if node.left.balance == 1:
node = self._rotation(node, left_rotation=False)
elif node.left.balance == -1:
node = self._double_rotation(node, left_side=True)
else:
# this state cannot accur
pass
else:
if node.right.balance == -1:
node = self._rotation(node, left_rotation=True)
elif node.right.balance == 1:
node = self._double_rotation(node, left_side=False)
else:
# this state cannot accur
pass
node.update()
return node
def delete(self, key) -> None:
self.root = self._delete(key, self.root)
def _delete(self, key, node: AvlNode) -> AvlNode:
if node.external:
return node
if node.key == key:
self._remove_node_from_chain(node)
if node.left == self.EXTERNAL_NODE:
return node.right
elif node.right == self.EXTERNAL_NODE:
return node.left
else:
replace_node = node.prev
node.left = self._delete(replace_node.key, node.left)
self._add_node_to_chain(replace_node, replace_node.prev, replace_node.nxt)
replace_node.left = node.left
replace_node.right = node.right
node = replace_node
elif node.key > key:
node.left = self._delete(key, node.left)
else:
node.right = self._delete(key, node.right)
left_depth = node.left.depth
right_depth = node.right.depth
if abs(left_depth - right_depth) <= 1:
# everything is OK
pass
else:
if left_depth > right_depth:
if node.left.balance == 1 or node.left.balance == 0:
node = self._rotation(node, left_rotation=False)
else:
node = self._double_rotation(node, left_side=True)
else:
if node.right.balance == -1 or node.right.balance == 0:
node = self._rotation(node, left_rotation=True)
else:
node = self._double_rotation(node, left_side=False)
node.update()
return node
def _rotation(self, subtree_root: AvlNode, left_rotation=True) -> AvlNode:
"""Makes single rotation, update node depth, return new root of subtree"""
if left_rotation:
right_node = subtree_root.right
subtree_root.right = right_node.left
right_node.left = subtree_root
subtree_root.update()
right_node.update()
return right_node
else:
left_node = subtree_root.left
subtree_root.left = left_node.right
left_node.right = subtree_root
subtree_root.update()
left_node.update()
return left_node
def _double_rotation(self, subtree_root: AvlNode, left_side=True) -> AvlNode:
"""Makes double rotation from single rotations, return new root of subtree"""
if left_side:
subtree_root.left = self._rotation(subtree_root.left)
return self._rotation(subtree_root, left_rotation=False)
else:
subtree_root.right = self._rotation(subtree_root.right, left_rotation=False)
return self._rotation(subtree_root)