Skip to content

Instantly share code, notes, and snippets.

@girish3
Created February 15, 2016 18:33
Show Gist options
  • Save girish3/a8e3931154af4da89995 to your computer and use it in GitHub Desktop.
Save girish3/a8e3931154af4da89995 to your computer and use it in GitHub Desktop.
AVL tree implementation in python
#import random, math
outputdebug = False
def debug(msg):
if outputdebug:
print msg
class Node():
def __init__(self, key):
self.key = key
self.left = None
self.right = None
class AVLTree():
def __init__(self, *args):
self.node = None
self.height = -1
self.balance = 0;
if len(args) == 1:
for i in args[0]:
self.insert(i)
def height(self):
if self.node:
return self.node.height
else:
return 0
def is_leaf(self):
return (self.height == 0)
def insert(self, key):
tree = self.node
newnode = Node(key)
if tree == None:
self.node = newnode
self.node.left = AVLTree()
self.node.right = AVLTree()
debug("Inserted key [" + str(key) + "]")
elif key < tree.key:
self.node.left.insert(key)
elif key > tree.key:
self.node.right.insert(key)
else:
debug("Key [" + str(key) + "] already in tree.")
self.rebalance()
def rebalance(self):
'''
Rebalance a particular (sub)tree
'''
# key inserted. Let's check if we're balanced
self.update_heights(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.lrotate() # we're in case II
self.update_heights()
self.update_balances()
self.rrotate()
self.update_heights()
self.update_balances()
if self.balance < -1:
if self.node.right.balance > 0:
self.node.right.rrotate() # we're in case III
self.update_heights()
self.update_balances()
self.lrotate()
self.update_heights()
self.update_balances()
def rrotate(self):
# Rotate left pivoting on self
debug ('Rotating ' + str(self.node.key) + ' right')
A = self.node
B = self.node.left.node
T = B.right.node
self.node = B
B.right.node = A
A.left.node = T
def lrotate(self):
# Rotate left pivoting on self
debug ('Rotating ' + str(self.node.key) + ' left')
A = self.node
B = self.node.right.node
T = B.left.node
self.node = B
B.left.node = A
A.right.node = T
def update_heights(self, recurse=True):
if not self.node == None:
if recurse:
if self.node.left != None:
self.node.left.update_heights()
if self.node.right != None:
self.node.right.update_heights()
self.height = max(self.node.left.height,
self.node.right.height) + 1
else:
self.height = -1
def update_balances(self, recurse=True):
if not self.node == None:
if recurse:
if self.node.left != None:
self.node.left.update_balances()
if self.node.right != None:
self.node.right.update_balances()
self.balance = self.node.left.height - self.node.right.height
else:
self.balance = 0
def delete(self, key):
# debug("Trying to delete at node: " + str(self.node.key))
if self.node != None:
if self.node.key == key:
debug("Deleting ... " + str(key))
if self.node.left.node == None and self.node.right.node == None:
self.node = None # leaves can be killed at will
# if only one subtree, take that
elif self.node.left.node == None:
self.node = self.node.right.node
elif self.node.right.node == None:
self.node = self.node.left.node
# worst-case: both children present. Find logical successor
else:
replacement = self.logical_successor(self.node)
if replacement != None: # sanity check
debug("Found replacement for " + str(key) + " -> " + str(replacement.key))
self.node.key = replacement.key
# replaced. Now delete the key from right child
self.node.right.delete(replacement.key)
self.rebalance()
return
elif key < self.node.key:
self.node.left.delete(key)
elif key > self.node.key:
self.node.right.delete(key)
self.rebalance()
else:
return
def logical_predecessor(self, node):
'''
Find the biggest valued node in LEFT child
'''
node = node.left.node
if node != None:
while node.right != None:
if node.right.node == None:
return node
else:
node = node.right.node
return node
def logical_successor(self, node):
'''
Find the smallese valued node in RIGHT child
'''
node = node.right.node
if node != None: # just a sanity check
while node.left != None:
debug("LS: traversing: " + str(node.key))
if node.left.node == None:
return node
else:
node = node.left.node
return node
def check_balanced(self):
if self == None or self.node == None:
return True
# We always need to make sure we are balanced
self.update_heights()
self.update_balances()
return ((abs(self.balance) < 2) and self.node.left.check_balanced() and self.node.right.check_balanced())
def inorder_traverse(self):
if self.node == None:
return []
inlist = []
l = self.node.left.inorder_traverse()
for i in l:
inlist.append(i)
inlist.append(self.node.key)
l = self.node.right.inorder_traverse()
for i in l:
inlist.append(i)
return inlist
def display(self, level=0, pref=''):
'''
Display the whole tree. Uses recursive def.
TODO: create a better display using breadth-first search
'''
self.update_heights() # Must update heights before balances
self.update_balances()
if(self.node != None):
print '-' * level * 2, pref, self.node.key, "[" + str(self.height) + ":" + str(self.balance) + "]", 'L' if self.is_leaf() else ' '
if self.node.left != None:
self.node.left.display(level + 1, '<')
if self.node.left != None:
self.node.right.display(level + 1, '>')
# Usage example
if __name__ == "__main__":
a = AVLTree()
print "----- Inserting -------"
#inlist = [5, 2, 12, -4, 3, 21, 19, 25]
inlist = [7, 5, 2, 6, 3, 4, 1, 8, 9, 0]
for i in inlist:
a.insert(i)
a.display()
print "----- Deleting -------"
a.delete(3)
a.delete(4)
# a.delete(5)
a.display()
print
print "Input :", inlist
print "deleting ... ", 3
print "deleting ... ", 4
print "Inorder traversal:", a.inorder_traverse()
@zinchse
Copy link

zinchse commented May 14, 2022

see my implementation :)
I couldn't find anything better
correct me if I'm wrong

@awwalm
Copy link

awwalm commented Jan 23, 2024

@zinchse : see my implementation :) I couldn't find anything better correct me if I'm wrong

+1
I can confirm this is the only implementation I have found that works. EXCELLENT job with the printer method, mate. Very clean and tidy too. Repo starred.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment