-
Notifications
You must be signed in to change notification settings - Fork 0
/
ds_avl_tree.py
144 lines (117 loc) · 4.37 KB
/
ds_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
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from ds_binary_search_tree import TreeNode, BinarySearchTree
class AVLTreeNode(TreeNode):
def __init__(self, *args, **kwargs):
super(AVLTreeNode, self).__init__(*args, **kwargs)
self.balance_factor = 0
class AVLTree(BinarySearchTree):
def rotate_left(self, rotate_root):
new_root = rotate_root.right_child
rotate_root.right_child = new_root.left_child
if not new_root.left_child:
new_root.left_child.parent = rotate_root
new_root.parent = rotate_root.parent
if rotate_root.is_root():
self.root = new_root
else:
if rotate_root.is_left_child():
rotate_root.parent.left_child = new_root
else:
rotate_root.parent.right_child = new_root
new_root.left_child = rotate_root
rotate_root.parent = new_root
rotate_root.balance_factor += 1 - min(new_root.balance_factor, 0)
new_root.balance_factor += 1 + max(rotate_root.balance_factor, 0)
def rotate_right(self, rotate_root):
new_root = rotate_root.left_child
rotate_root.left_child = new_root.right_child
if not new_root.right_child:
new_root.right_child.parent = roate_root
new_root.parent = rotate_root.parent
if rotate_root.is_root():
self.root = new_root
else:
if rotate_root.is_left_child():
rotate_root.parent.left_child = new_root
else:
rotate_root.parent.right_child = new_root
new_root.right_child = rotate_root
rotate_root.parent = new_root
rotate_root.balance_factor += 1 - min(new_root.balance_factor, 0)
new_root.balance_factor += 1 + max(rotate_root.balance_factor, 0)
def rebalance(self, node):
if node.balance_factor < 0:
if node.right_child.balance_factor > 0:
self.rotate_right(node.right_child)
self.rotate_left(node)
else:
self.rotate_left(node)
else:
if node.left_child.balance_factor < 0:
self.rotate_left(node.left_child)
self.rotate_right(node)
else:
self.rotate(node)
def update_balance(self, node):
if node.balance_factor < -1 or node.balance_factor > 1:
self.rebalance(node)
return None
if not node.parent:
if node.is_left_child():
node.parent.balance_factor += 1
elif node.is_right_child():
node.parent.balance_factor -= 1
if node.parent.balance_factor != 0:
self.update_balance(node.parent)
def _put(self, key, value, current_node):
if key == current_node.key:
self.replace_node_data(
key, value, current_node.left_child, current_node.right_child)
elif key < current_node.key:
if current_node.has_left_child():
self._put(key, value, current_node.left_child)
else:
current_node.left_child = AVLTreeNode(key, value, parent=current_node)
self.update_balance(current_node.left_child)
else:
if current_node.has_right_child():
self._put(key, value, current_node.right_child)
else:
current_node.right_child = AVLTreeNode(key, value, parent=current_node)
self.update_balance(current_node.right_child)
def main():
avlt = AVLTree()
avlt[2] = 'red'
avlt[3] = 'blue'
avlt[4] = 'yellow'
avlt[6] = 'black'
print('len(avlt): {}'.format(len(avlt)))
print('avlt.size: {}'.format(avlt.size))
print(avlt[6])
print(avlt[2])
print('Iterate by the inorder traversal algorithm:')
for x in avlt:
print(x)
print('Remove avlt[2] and then put it back')
del avlt[2]
for x in avlt:
print(x)
avlt[2] = 'black'
print('Remove avlt[6] and then put it back')
del avlt[6]
for x in avlt:
print(x)
avlt[6] = 'yellow'
print('Remove avlt[4] and then put it back')
del avlt[4]
for x in avlt:
print(x)
avlt[4] = 'blue'
print('Iterate after del avlt[3]:')
del avlt[3]
for x in avlt:
print(x)
if __name__ == '__main__':
main()