Skip to content

Instantly share code, notes, and snippets.

@SubCoder1
Created January 12, 2020 22:03
Show Gist options
  • Save SubCoder1/a585797d8e481dae8c62916e6fe5464b to your computer and use it in GitHub Desktop.
Save SubCoder1/a585797d8e481dae8c62916e6fe5464b to your computer and use it in GitHub Desktop.
AVL Tree implementation in Python3
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