Created
April 7, 2013 01:13
-
-
Save rcloran/5328406 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
class Empty: | |
pass | |
class AVLTree(object): | |
value = None | |
left = None | |
right = None | |
def __init__(self, value=Empty, left=None, right=None, parent=None): | |
if value is Empty and parent: | |
raise ValueError("Cannot use special value Empty") | |
self.value = value | |
self.left = left | |
self.right = right | |
self.parent = parent | |
@property | |
def height(self): | |
if self.left: | |
lh = self.left.height | |
else: | |
lh = 0 | |
if self.right: | |
rh = self.right.height | |
else: | |
rh = 0 | |
return max(lh, rh) + 1 | |
def search(self, value, mutator=None): | |
"""Return the sub-tree which contains value value""" | |
if mutator is None: | |
mutator = lambda x: x | |
if value == mutator(self.value): | |
return self | |
if value < self.value: | |
if self.left: | |
return self.left.search(value) | |
return self # self contains -infinity to right, ie value | |
if value > self.value: | |
if self.right: | |
return self.right.search(value) | |
return self | |
def insert(self, value): | |
if value is Empty: | |
raise ValueError("Cannot insert special value Empty") | |
if self.value is Empty: | |
self.value = value | |
return | |
if value <= self.value: | |
if self.left: | |
return self.left.insert(value) | |
else: | |
self.left = AVLTree(value, parent=self) | |
self.rebalance() | |
return | |
if self.right: | |
self.right.insert(value) | |
else: | |
self.right = AVLTree(value, parent=self) | |
self.rebalance() | |
def insert_tree(self, other): | |
for value in other: | |
self.insert(value) | |
def __iter__(self): | |
if self.left: | |
for value in self.left: | |
yield value | |
if self.value is not Empty: | |
yield self.value | |
if self.right: | |
for value in self.right: | |
yield value | |
def _replace_with(self, other): | |
self.left = other.left | |
self.value = other.value | |
self.right = other.right | |
def remove_child(self, other): | |
if other is self.left: | |
self.left = None | |
elif other is self.right: | |
self.right = None | |
else: | |
raise ValueError('whut') | |
def delete(self, value): | |
if value != self.value: | |
tree = self.search(value) | |
if tree.value != value: | |
raise KeyError("Value not in tree") | |
return tree.delete(value) | |
if not self.left and not self.right: | |
if not self.parent: | |
self.value = Empty | |
return | |
self.parent.remove_child(self) | |
return | |
if not self.left: | |
return self._replace_with(self.right) | |
if not self.right: | |
return self._replace_with(self.left) | |
right_min = self.right.min_node() | |
self.value = right_min.value | |
# The following is basically right_min.delete(value) | |
if right_min.right: | |
right_min.value = right_min.right.value | |
right_min.right = None | |
elif right_min.parent is not self: | |
right_min.parent.left = None | |
else: # We've deleted the last remaining right node | |
self.right = None | |
self.rebalance() | |
def rebalance(self): | |
bf = self.balance_factor() | |
dh = abs(bf) | |
if dh <= 1: | |
if self.parent: | |
return self.parent.rebalance() | |
return | |
if bf > 0: # Left heavy | |
left_bf = self.left.balance_factor() | |
if left_bf >= 0: | |
self.rotate_right() | |
return | |
else: | |
self.left.rotate_left() | |
self.rotate_right() | |
else: # right heavy | |
right_bf = self.right.balance_factor() | |
if right_bf < 0: | |
self.rotate_left() | |
return | |
else: | |
self.right.rotate_right() | |
self.rotate_left() | |
if self.parent: | |
return self.parent.rebalance() | |
return | |
def rotate_right(self): | |
if not self.left: | |
raise ValueError("Cannot rotate right") | |
a, b, g = self.left.left, self.left.right, self.right | |
self.right = AVLTree(self.value, left=b, right=g, parent=self) | |
if b: | |
b.parent = self.right | |
if g: | |
g.parent = self.right | |
self.value = self.left.value | |
self.left = a | |
if self.left: | |
self.left.parent = self | |
def rotate_left(self): | |
if not self.right: | |
raise ValueError("Cannot rotate left") | |
a, b, g = self.left, self.right.left, self.right.right | |
self.left = AVLTree(self.value, left=a, right=b, parent=self) | |
if a: | |
a.parent = self.left | |
if b: | |
b.parent = self.left | |
self.value = self.right.value | |
self.right = g | |
if self.right: | |
self.right.parent = self | |
def printtree(self, level=0): | |
if self.right: | |
self.right.printtree(level + 1) | |
print((" " * level) + str(self.value)) | |
if self.left: | |
self.left.printtree(level + 1) | |
def min_node(self): | |
if not self.left: | |
return self | |
return self.left.min_node() | |
def min(self): | |
return self.min_node().value | |
def max_node(self): | |
if not self.right: | |
return self | |
return self.right.max_node() | |
def max(self): | |
return self.max_node().value | |
def balance_factor(self): | |
if not self.left: | |
lh = 0 | |
else: | |
lh = self.left.height | |
if not self.right: | |
rh = 0 | |
else: | |
rh = self.right.height | |
return lh - rh | |
if __name__ == '__main__': | |
t = AVLTree(5) | |
t.printtree() | |
t.insert(3) | |
t.insert(6) | |
t.printtree() | |
t.insert(4) | |
t.insert(2) | |
t.insert(1) | |
t.insert(0) | |
t.insert(-1) | |
t.printtree() | |
t.delete(-1) | |
t.delete(1) | |
t.printtree() | |
t.delete(5) | |
t.printtree() | |
t.delete(4) | |
t.printtree() | |
t.delete(0) | |
t.delete(2) | |
t.delete(3) | |
t.printtree() | |
t.delete(6) | |
t.printtree() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment