Created
February 15, 2016 18:33
-
-
Save girish3/a8e3931154af4da89995 to your computer and use it in GitHub Desktop.
AVL tree implementation in python
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
| #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 "Input :", inlist | |
| print "deleting ... ", 3 | |
| print "deleting ... ", 4 | |
| print "Inorder traversal:", a.inorder_traverse() |
I wonder this code can't work well as root is balanced while sub tree is not
How come the function logical_predecessor is never used?
How can I search a key?
see my implementation :)
I couldn't find anything better
correct me if I'm wrong
@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
I guess the code is taken from here.
https://blog.coder.si/2014/02/how-to-implement-avl-tree-in-python.html