1+ # AVL Tree
2+ # On every insertion check if tree is balanced or not using height
3+ # If it is unbalanced, make rotation to balance it
4+ 5+ 6+ class Node (object ):
7+ 8+ def __init__ (self ,data ):
9+ self .data = data
10+ self .leftChild = None # Left Child to current node
11+ self .rightChild = None # Right cchild to current node
12+ self .height = 0 # Height of Tree to make sure tree is balanced
13+ 14+ 15+ # Perform operations on AVL Tree
16+ class AVL (object ):
17+ def __init__ (self ):
18+ self .root = None # Root Node
19+ 20+ # --------------------------- Calculate Height of AVL Tree ---------------------------------
21+ def calcHeight (self ,node ):
22+ if not node : # If this Node is Null
23+ return - 1 # return -1; left and right child of leaf Nodes
24+ #print('\nHeight: ', node.height)
25+ return node .height
26+ 27+ # ------------------------------- Insertion of Node in a AVL Tree ----------------------------
28+ def insert (self , data ):
29+ self .root = self .insertNode (data , self .root )
30+ 31+ def insertNode (self , data , node ):
32+ if not node : # If it is a root node, create a new node
33+ return Node (data )
34+ if data < node .data : # If data < current node, go to left else right
35+ node .leftChild = self .insertNode (data , node .leftChild )
36+ else :
37+ node .rightChild = self .insertNode (data , node .rightChild )
38+ # Get the height of the new node
39+ node .height = max (self .calcHeight (node .leftChild ), self .calcHeight (node .rightChild )) + 1
40+ # Node has been inserted.
41+ # Now check if there has been any violations of AVL Tree
42+ #print('Node {} Inserted'.format(data))
43+ return self .settleViolation (data , node )
44+ 45+ 46+ # -------------------- Settle any Violations upon adding/inserting new Node --------------------------
47+ def settleViolation (self , data , node ):
48+ balance = self .calcBalance (node )
49+ # Case 1: Left Left Heavy situation
50+ if balance > 1 and data < node .leftChild .data :
51+ print ("Left Left Heavy Situation..." )
52+ return self .rotateRight (node )
53+ # Case 2: Right Right Heavy Situation
54+ if balance < - 1 and data > node .rightChild .data :
55+ print ('Right Right Heavy Situation...' )
56+ return self .rotateLeft (node )
57+ # Case 3: Left Right Heavy Situation
58+ if balance > 1 and data > node .leftChild .data :
59+ print ('Left Right Heavy Situation...' )
60+ node .leftChild = self .rotateLeft (node .leftChild )
61+ return self .rotateRight (node )
62+ # Case 4
63+ if balance < - 1 and data < node .rightChild .data :
64+ print ('Right Left Heavy Situation...' )
65+ node .rightChild = self .rotateRight (node .rightChild ) # Here node is the Root Node
66+ return self .rotateLeft (node )
67+ return node
68+ 69+ 70+ # ------------------------------- See if Tree is Balanced or not -------------------------------
71+ # If it returns value > 1, it means it is a Left heavy tree
72+ # Make Right rotation to balance it
73+ # If it returns value < 1, it means it is a Right heavy tree
74+ # Make Left rotation to balance it
75+ def calcBalance (self ,node ):
76+ if not node :
77+ return 0
78+ return self .calcHeight (node .leftChild ) - self .calcHeight (node .rightChild )
79+ 80+ # Rotate Nodes to Right to Balance AVL Tree
81+ # O(1) time Complexity
82+ def rotateRight (self ,node ):
83+ print ('Rotating to right on node ' , node .data ) # C <- B <- D -> E; Root node is "D"
84+ tempLeftChild = node .leftChild # tempLeftChild => B
85+ t = tempLeftChild .rightChild # t = C
86+ 87+ # Rotate Right
88+ tempLeftChild .rightChild = node # "D" becomes right child of "B"; B -> D
89+ node .leftChild = t # "C" becomes left child of "D"
90+ 91+ # Calculate Height of AVL Tree
92+ node .height = max (self .calcHeight (node .leftChild ), self .calcHeight (node .rightChild )) + 1
93+ tempLeftChild .height = max (self .calcHeight (tempLeftChild .leftChild ), self .calcHeight (tempLeftChild .rightChild )) + 1
94+ return tempLeftChild # Root Node after Rotation
95+ 96+ 97+ # Rotate Nodes to Left to Balance AVL Tree
98+ # O(1) time Complexity
99+ def rotateLeft (self ,node ):
100+ print ('Rotating to Left on node ' , node .data )
101+ tempRightChild = node .rightChild
102+ t = tempRightChild .leftChild
103+ 104+ # Rotate Right
105+ tempRightChild .leftChild = node
106+ node .rightChild = t
107+ 108+ # Calculate Height of AVL Tree
109+ node .height = max (self .calcHeight (node .leftChild ), self .calcHeight (node .rightChild )) + 1
110+ tempRightChild .height = max (self .calcHeight (tempRightChild .leftChild ), self .calcHeight (tempRightChild .rightChild )) + 1
111+ return tempRightChild # Root Node after Rotation
112+ 113+ 114+ # ------------------------------ Remove Node from AVL Tree -----------------------
115+ # Remove Node [Deletion]
116+ def remove (self , data ):
117+ if self .root :
118+ self .root = self .removeNode (data , self .root )
119+ 120+ def removeNode (self , data , node ):
121+ if not node :
122+ return node
123+ if data < node .data :
124+ node .leftChild = self .removeNode (data ,node .leftChild )
125+ if data > node .data :
126+ node .rightChild = self .removeNode (data , node .rightChild )
127+ else :
128+ if not node .leftChild and not node .rightChild :
129+ print ('Removing a leaf node...' )
130+ del node
131+ return None
132+ if not node .leftChild :
133+ print ('Removing right child...' )
134+ tempNode = node .rightChild
135+ del node
136+ return tempNode
137+ if not node .rightChild :
138+ print ('Removing left child...' )
139+ tempNode = node .leftChild
140+ return tempNode
141+ print ('Removing Node with two children...' )
142+ tempNode = self .getPredecessor (node .leftChild )
143+ node .data = tempNode .data
144+ node .leftChild = self .removeNode (tempNode .data , node .leftChild )
145+ if not node :
146+ return node
147+ 148+ node .height = max (self .calcHeight (node .leftChild ), self .calcHeight (node .rightChild )) + 1
149+ balance = self .calcBalance (node )
150+ 151+ # Doubly Left Heavy Tree
152+ if balance > 1 and self .calcBalance (node .leftChild ) >= 0 :
153+ return self .rotateRight (node )
154+ # Doubly Right Heavy Tree
155+ if balance < - 1 and self .calcBalance (node .rightChild ) <= 0 :
156+ return self .rotateLeft (node )
157+ # Left Right Case
158+ if balance > 1 and self .calcBalance (node .leftChild ) < 0 :
159+ node .leftChild = self .rotateLeft (node .leftChild )
160+ return self .rotateRight (node )
161+ # Right Left Case
162+ if balance < - 1 and self .calcBalance (node .rightChild ) > 0 :
163+ node .rightChild = self .rotateRight (node .rightChild )
164+ return self .rotateLeft (node )
165+ return node
166+ 167+ 168+ def getPredecessor (self , node ):
169+ if node .rightChild :
170+ return self .getPredecessor (node .rightChild )
171+ return node
172+ 173+ 174+ # --------------------- Traverse the AVL Tree ------------------
175+ def traverse (self ):
176+ if self .root :
177+ self .traverseInOrder (self .root )
178+ 179+ 180+ def traverseInOrder (self , node ):
181+ if node .leftChild :
182+ self .traverseInOrder (node .leftChild )
183+ print ('%s' % node .data )
184+ if node .rightChild :
185+ self .traverseInOrder (node .rightChild )
186+ 187+ # ------------------- Testing -----------------
188+ avl = AVL ()
189+ avl .insert (10 )
190+ avl .insert (20 )
191+ avl .insert (5 )
192+ avl .insert (6 )
193+ avl .insert (15 )
194+ avl .traverse ()
195+ 196+ avl .remove (20 )
197+ avl .remove (15 )
198+ avl .traverse ()
199+ 200+ #---------------------------------------- EOC ------------------------------------
0 commit comments