-
Notifications
You must be signed in to change notification settings - Fork 0
/
avltree.py
158 lines (145 loc) · 5.5 KB
/
avltree.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
153
154
155
156
157
158
from binarysearchtree import BSTNode, BinarySearchTree
class AVLNode(BSTNode):
def __init__(self, key, value = None):
BSTNode.__init__(self, key, value)
self.balanceFactor = 0
class AVLTree(BinarySearchTree):
def __init__(self):
BinarySearchTree.__init__(self)
def insert(self, key, value = None):
node = AVLNode(key, value)
if self.root:
self._insert(node, self.root)
else:
self.root = node
self.count += 1
def _insert(self, node, currentNode):
if node.key < currentNode.key:
if currentNode.leftChild:
self._insert(node, currentNode.leftChild)
else:
node.parent = currentNode
currentNode.leftChild = node
self._balance(node)
elif node.key > currentNode.key:
if currentNode.rightChild:
self._insert(node, currentNode.rightChild)
else:
node.parent = currentNode
currentNode.rightChild = node
self._balance(node)
else:
currentNode.value = node.value
def _balance(self, node):
if node.balanceFactor > 1 or node.balanceFactor < -1:
self._rebalance(node)
return
if node.parent != None:
if node == node.parent.leftChild:
node.parent.balanceFactor += 1
else:
node.parent.balanceFactor -= 1
if node.parent.balanceFactor != 0:
self._balance(node.parent)
def _rebalance(self, node):
if node.balanceFactor < 0:
if node.rightChild.balanceFactor > 0:
self._rotateRight(node.rightChild)
self._rotateLeft(node)
else:
self._rotateLeft(node)
elif node.balanceFactor > 0:
if node.leftChild.balanceFactor < 0:
self._rotateLeft(node.leftChild)
self._rotateRight(node)
else:
self._rotateRight(node)
def _rotateLeft(self, node):
newRoot = node.rightChild
node.rightChild = newRoot.leftChild
if newRoot.leftChild != None:
newRoot.leftChild.parent = node
newRoot.parent = node.parent
if node.parent == None:
self.root = newRoot
else:
if node == node.parent.leftChild:
node.parent.leftChild = newRoot
else:
node.parent.rightChild = newRoot
newRoot.leftChild = node
node.parent = newRoot
node.balanceFactor = node.balanceFactor + \
1 - min(newRoot.balanceFactor, 0)
newRoot.balanceFactor = newRoot.balanceFactor + \
1 + max(node.balanceFactor, 0)
def _rotateRight(self, node):
newRoot = node.leftChild
node.leftChild = newRoot.rightChild
if newRoot.rightChild != None:
newRoot.rightChild.parent = node
newRoot.parent = node.parent
if node.parent == None:
self.root = newRoot
else:
if node == node.parent.rightChild:
node.parent.rightChild = newRoot
else:
node.parent.leftChild = newRoot
newRoot.rightChild = node
node.parent = newRoot
node.balanceFactor = node.balanceFactor + \
1 + max(newRoot.balanceFactor, 0)
newRoot.balanceFactor = newRoot.balanceFactor + \
1 - min(node.balanceFactor, 0)
def __delitem__(self, key):
self.delete(key)
def delete(self, key):
if self.count == 1:
if self.root.key == key:
self.root = None
self.count -= 1
return
elif self.count > 1:
node = self._find(key, self.root)
if node:
self._delete(node)
return
raise KeyError('{0} not found.'.format(key))
def _delete(self, node):
# handle 'neither', 'left', 'right', and 'both' children cases
if node.leftChild is None and node.rightChild is None:
if node is node.parent.leftChild:
node.parent.leftChild = None
self._balance(node.parent)
else:
node.parent.rightChild = None
self._balance(node.parent)
self.count -= 1
elif node.leftChild and node.rightChild is None:
if node.parent is None:
self.root = node.leftChild
elif node is node.parent.leftChild:
node.parent.leftChild = node.leftChild
self._balance(node.parent)
else:
node.parent.rightChild = node.leftChild
self._balance(node.parent)
node.leftChild.parent = node.parent
self.count -= 1
elif node.rightChild and node.leftChild is None:
if node.parent is None:
self.root = node.rightChild
elif node is node.parent.leftChild:
node.parent.leftChild = node.rightChild
self._balance(node.parent)
else:
node.parent.rightChild = node.rightChild
self._balance(node.parent)
node.rightChild.parent = node.parent
self.count -= 1
else:
successor = node.find_successor()
self._delete(successor)
node.key = successor.key
node.value = successor.value