girish3 / avl_tree.py
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 () |
Источник