Created
January 12, 2020 22:03
-
-
Save SubCoder1/a585797d8e481dae8c62916e6fe5464b to your computer and use it in GitHub Desktop.
AVL Tree implementation in Python3
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
class Node: | |
def __init__(self, val=0): | |
self.left = self.right = self.parent = None | |
self.height = 0 | |
self.val = val | |
class AVL: | |
def __init__(self): | |
self.root = None | |
def InsertNode(self, val): | |
if self.root is None: | |
self.root = Node(val=val) | |
else: | |
node_ptr = self.root | |
node = Node(val=val) | |
while node_ptr is not None: | |
if node_ptr.val > node.val: | |
if node_ptr.left is None: | |
node_ptr.left = node | |
node.parent = node_ptr | |
break | |
else: | |
node_ptr = node_ptr.left | |
else: | |
if node_ptr.right is None: | |
node_ptr.right = node | |
node.parent = node_ptr | |
break | |
else: | |
node_ptr = node_ptr.right | |
self.BalanceTree(node) | |
def getHeight(self, node): | |
return -1 if node is None else node.height | |
def setHeight(self, node): | |
l = node.left.height if node.left is not None else -1 | |
r = node.right.height if node.right is not None else -1 | |
node.height = max(l,r) + 1 | |
def BalanceFactor(self, node): | |
return (self.getHeight(node.right) - self.getHeight(node.left)) | |
def BalanceTree(self, node): | |
if node is self.root: | |
self.Treefix(self.root) | |
else: | |
while node is not None: | |
self.setHeight(node) | |
node = node.parent | |
self.Treefix(node) | |
def Treefix(self, node): | |
if node is not None: | |
if self.BalanceFactor(node) is 2: | |
if self.BalanceFactor(node.right) < 0: | |
self.RightRotate(node.right) | |
self.LeftRotate(node) | |
self.setHeight(node) | |
elif self.BalanceFactor(node) is -2: | |
if self.BalanceFactor(node.left) > 0: | |
self.LeftRotate(node.left) | |
self.RightRotate(node) | |
self.setHeight(node) | |
def LeftRotate(self, node): | |
tmp_node = Node() | |
if node.right.left is not None: | |
tmp_node.right = node.right.left | |
tmp_node.right.parent = tmp_node | |
if node.left is not None: | |
tmp_node.left = node.left | |
tmp_node.left.parent = tmp_node | |
tmp_node.val = node.val | |
node.val = node.right.val | |
node.left = tmp_node | |
self.setHeight(tmp_node) | |
if node.right.right is not None: | |
node.right = node.right.right | |
node.right.parent = node | |
else: | |
node.right = None | |
self.setHeight(node.right) | |
self.setHeight(node) | |
def RightRotate(self, node): | |
tmp_node = Node() | |
if node.left.right is not None: | |
tmp_node.left = node.left.right | |
tmp_node.left.parent = tmp_node | |
if node.right is not None: | |
tmp_node.right = node.right | |
tmp_node.right.parent = tmp_node | |
tmp_node.val = node.val | |
node.val = node.left.val | |
node.right = tmp_node | |
tmp_node.parent = node | |
self.setHeight(tmp_node) | |
if node.left.left is not None: | |
node.left = node.left.left | |
node.left.parent = node | |
else: | |
node.left = None | |
self.setHeight(node.left) | |
self.setHeight(node) | |
def InorderTraversal(self, node): | |
if node is not None: | |
self.InorderTraversal(node.left) | |
print(node.val, end=' ') | |
self.InorderTraversal(node.right) | |
if __name__ == '__main__': | |
bst = AVL() | |
bst.InsertNode(10) | |
bst.InsertNode(20) | |
bst.InsertNode(5) | |
bst.InsertNode(4) | |
bst.InsertNode(3) | |
bst.InorderTraversal(bst.root) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment