https://www.geeksforgeeks.org/avl-tree-set-1-insertion/

AVL Tree : Insertion.py

#!/usr/bin/env python3


class Node(object):
	"""Class for tree nodes"""

	def __init__(self, key):
		super(Node, self).__init__()
		self.key = key
		self.left = None
		self.right = None
		self.height = 1


def preorder(root):
	if root is None:
		return

	print("{}".format(root.key), end=" ")
	preorder(root.left)
	preorder(root.right)


def inorder(root):
	if root is None:
		return

	inorder(root.left)
	print("{}".format(root.key), end=" ")
	inorder(root.right)


def postorder(root):
	if root is None:
		return

	postorder(root.left)
	postorder(root.right)
	print("{}".format(root.key), end=" ")


class AVL(object):
	"""Class for implementation of AVL tree"""

	def __init__(self):
		super(AVL, self).__init__()

	def insert(self, root, key):
		# insert the node like a normal BST
		if root is None:
			return Node(key)
		elif key < root.key:
			root.left = self.insert(root.left, key)
		else:  # key > root.key
			root.right = self.insert(root.right, key)

		# update the height of the parent node
		root.height = self.update_height(root)

		# get the balance factor
		balance_factor = self.get_balance(root)

		"""
		These conditions only work when there are three nodes
		
		For right-left:
			key = 6
			root = 5
			root.right.key = 7
		
		Therefore, 
			key < root.right.key 
			
				5
				 \
				  7
				 /
				6
								
		For right-right:
			key = 7
			root = 5
			root.right.key = 6
			
		Therefore,
			key > root.right.key
			
				5
				 \
				  6
				   \
				    7
				    
		Similarly, for left-left and left-right
		
		"""

		# right heavy
		if balance_factor > 1:
			# right right
			if key > root.right.key:
				return self.left_rotate(root)
			# right left
			elif key < root.right.key:
				root.right = self.right_rotate(root.right)
				return self.left_rotate(root)

		# left heavy
		elif balance_factor < -1:
			# left left
			if key < root.left.key:
				return self.right_rotate(root)
			# left right
			elif key < root.left.key:
				root.left = self.left_rotate(root.left)
				return self.right_rotate(root)

		return root

	"""
	T1, T2, T3 and T4 are subtrees
	
		Left rotation:
		
		   z                                  y
		  / \                               /   \ 
		T1   y        Left Rotate(z)       z     x
		    /  \      - - - - - - - ->    / \   / \
		   T2   x                        T1 T2 T3 T4
		       / \
		     T3  T4
		
		
		Right rotation:
		
		 z                                  y 
	        / \                               /   \
	       y   T4     Right Rotate (z)       x     z
	      / \         - - - - - - - ->      / \   / \ 
	     x   T3                            T1 T2 T3 T4
	    / \
	   T1 T2
	  
	"""

	def left_rotate(self, z):
		y = z.right
		t2 = y.left

		# performing rotation
		y.left = z
		z.right = t2

		# update heights
		z.height = self.update_height(z)
		y.height = self.update_height(y)

		# return new root
		return y

	def right_rotate(self, z):
		y = z.left
		t3 = y.right

		# performing rotation
		y.right = z
		z.left = t3

		# update heights
		z.height = self.update_height(z)
		y.height = self.update_height(y)

		# return new root
		return y

	@staticmethod
	def get_height(root):
		if root is None:
			return 0
		else:
			return root.height

	def update_height(self, root):
		if root is None:
			return 0
		else:
			return 1 + max(self.get_height(root.right), self.get_height(root.left))

	def get_balance(self, root):
		if root is None:
			return 0
		else:
			return self.get_height(root.right) - self.get_height(root.left)


def main():
	tree = AVL()
	root = None

	root = tree.insert(root, 10)
	root = tree.insert(root, 20)
	root = tree.insert(root, 30)
	root = tree.insert(root, 40)
	root = tree.insert(root, 50)
	root = tree.insert(root, 25)

	"""
	The constructed AVL Tree would be
	            30
	           /  \
	         20    40
	        /  \    \
	       10  25    50
	"""

	print("Preorder traversal of the created tree is")
	preorder(root)
	print("\n")

	print("Inorder traversal of the created tree is")
	inorder(root)
	print("\n")

	print("Postorder traversal of the created tree is")
	postorder(root)
	print("\n")


if __name__ == '__main__':
	main()