Skip to content

Instantly share code, notes, and snippets.

@Twoody
Forked from girish3/avl_tree.py
Created October 28, 2018 20:03
Show Gist options
  • Save Twoody/de8d079842e0dd20cf20d870c73168af to your computer and use it in GitHub Desktop.
Save Twoody/de8d079842e0dd20cf20d870c73168af 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())
@swhsiang
Copy link

swhsiang commented Jan 6, 2019

@adeshsangoi the minimum of tree would be the left most node in the tree (it's a binary search tree).

@Xinergy
Copy link

Xinergy commented Dec 9, 2019

Line 229 should be if self.node.right != None: ?

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