forked from keon/algorithms
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathavl.py
More file actions
124 lines (105 loc) · 3.33 KB
/
Copy pathavl.py
File metadata and controls
124 lines (105 loc) · 3.33 KB
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
from tree.tree import TreeNode
class AvlTree(object):
"""
An avl tree.
"""
def __init__(self):
# Root node of the tree.
self.node = None
self.height = -1
self.balance = 0
def insert(self, key):
"""
Insert new key into node
"""
# Create new node
n = TreeNode(key)
if not self.node:
self.node = n
self.node.left = AvlTree()
self.node.right = AvlTree()
elif key < self.node.val:
self.node.left.insert(key)
elif key > self.node.val:
self.node.right.insert(key)
self.re_balance()
def re_balance(self):
"""
Re balance tree. After inserting or deleting a node,
"""
self.update_heights(recursive=False)
self.update_balances(False)
while self.balance < -1 or self.balance > 1:
if self.balance > 1:
if self.node.left.balance < 0:
self.node.left.rotate_left()
self.update_heights()
self.update_balances()
self.rotate_right()
self.update_heights()
self.update_balances()
if self.balance < -1:
if self.node.right.balance > 0:
self.node.right.rotate_right()
self.update_heights()
self.update_balances()
self.rotate_left()
self.update_heights()
self.update_balances()
def update_heights(self, recursive=True):
"""
Update tree height
"""
if self.node:
if recursive:
if self.node.left:
self.node.left.update_heights()
if self.node.right:
self.node.right.update_heights()
self.height = 1 + max(self.node.left.height, self.node.right.height)
else:
self.height = -1
def update_balances(self, recursive=True):
"""
Calculate tree balance factor
"""
if self.node:
if recursive:
if self.node.left:
self.node.left.update_balances()
if self.node.right:
self.node.right.update_balances()
self.balance = self.node.left.height - self.node.right.height
else:
self.balance = 0
def rotate_right(self):
"""
Right rotation
"""
new_root = self.node.left.node
new_left_sub = new_root.right.node
old_root = self.node
self.node = new_root
old_root.left.node = new_left_sub
new_root.right.node = old_root
def rotate_left(self):
"""
Left rotation
"""
new_root = self.node.right.node
new_left_sub = new_root.left.node
old_root = self.node
self.node = new_root
old_root.right.node = new_left_sub
new_root.left.node = old_root
def in_order_traverse(self):
"""
In-order traversal of the tree
"""
result = []
if not self.node:
return result
result.extend(self.node.left.in_order_traverse())
result.append(self.node.key)
result.extend(self.node.right.in_order_traverse())
return result