# AVL Tree implementation Based on VisualGo

I have tried implementing an AVL Tree on my own, based on visualising it. But i'm unsure how many testcases it work with, and how efficient it is. Are there any ways to make it efficient, and compact?

class Node:
def __init__(self, value, parent=None):
"""
Every Node has a value, and a height based on how far it is from the root node
"""
self.value = value
self.parent = parent
self.left = None
self.right = None

class AVLTree:
def __init__(self, List=None):
"""
initialising the root which is the first node which is None in an empty Tree

If List keyword argument is provided, then it will insert every value in the List.
"""
self.root = None
if List:
for item in List:
self.insert(item)

def inorder(self, current_node):
"""
Recursive Inorder Traversal of the Binary Tree, to display from smallest to largest, the values in the tree
Ex:
1 2 3 5 6 9
"""
if current_node:
self.inorder(current_node.left)
print(current_node.value, end=' ')
self.inorder(current_node.right)

def get_depth(self, node):
"""
Finds the depth of a specific node, which means the levels of nodes beneath it, also recursive.
Ex:                               height:
O                       0
/        \
O          O                 1
/     \    /      \
O       O  O        O            2
/
O              3

The Depth of root node (height 0) = 3, since there are 3 levels of nodes beneath it
The Depth of left node at (height 1) = 1, since there is only 1 level of nodes beneath it
The Depth of right node at (height 1) = 2, since there are   2                
The Depth of left node of (height 2) = 0, since there are no nodes beneath it
The Depth of a Non Existent Node = 0
"""
# Null node has 0 depth.
if node is None:
return 0

# Get the depth of the left and right subtree
# using recursion.
leftDepth = self.get_depth(node.left)
rightDepth = self.get_depth(node.right)

# Choose the larger one and add the root to it.
if leftDepth > rightDepth:
return leftDepth + 1
elif leftDepth < rightDepth:
return rightDepth + 1
else: # A Node with no children has a depth of 0
return 0

def get_balance_factor(self, node):
"""
Balance_factor is the difference between the heights of the two child node of the node

Rebalancing occurs when the balance_factor is greater than a difference of 1 (or) less than a difference of -1

usage: balance_factor = get_balance_factor(node)
which is left_child_height - right_child_height

If balanced, balance_factor = -1/0/1:
O           O          O         O
/ \         / \        / \       / \
O   O       O   O      O   O     O   O
/ \        /               \
O   O      O                 O

If balance_factor > 1:
O    <--- Node
/ \
O   O   Height = 2
/ \
O   O     Height = 3
/ \
O   O       Height = 4
Difference of L and R = 2 (4-2)

If balance_factor < 1:
O   <--- Node
/ \
O   O        Height = 2
/ \
O   O      Height = 3
/ \
O   O    Height = 4
Difference of L and R = -2 (2-4)

"""
if node is None:
return
return self.get_depth(node.left) - self.get_depth(node.right)

def rotate_right(self, y):
"""
Rotates right
Ex:

y                               x
/ \     Right Rotation          /  \
x   T3   - - - - - - - >        T1   y
/ \                                  / \
T1  T2                              T2  T3

changed:
1) the left child of y became T2
2) the right child of x became y
3) y's parent now points to x
4) T2's parent is now y if it exists
5) x's parent is now y's parent

"""
# Rotation
parent = y.parent

x = y.left
T2 = x.right

x.right = y
y.left = T2

if T2:
T2.parent = y
x.parent = parent

# If there is no parent, then that means y was the root node
if parent is None:
self.root = x
elif parent.left == y:
parent.left = x
elif parent.right == y:
parent.right = x

def rotate_left(self, x):
"""
Rotates left
Ex:

y                               x
/ \                             /  \
x   T3                          T1   y
/ \       < - - - - - - -            / \
T1  T2     Left Rotation            T2  T3

changed:
1) the left child of y became x
2) the right child of x became T2
3) x's parent now points to y
4) T2's parent is now x if it exists
5) y's parent is now x's parent

"""
# Rotation
parent = x.parent
y = x.right
T2 = y.left

y.left = x
x.right = T2

if T2:
T2.parent = x
y.parent = parent

# If there is no parent, then that means y was the root node
if parent is None:
self.root = y
elif parent.left == x:
parent.left = y
elif y.parent.right == x:
parent.right = y

def recursive_balancing_for_insertion(self, current_node, value):
"""
As the name says, rebalances, by using comparisons between the value to be inserted and child node's value
It slowly checks the balance factor of every node from the parent of the recently inserted node,
after a node has been checked, it goes the parent of that node, until it reaches the root node, which doesn't have a parent /= None.

Depending on the balance factor and comparison, it may perform a left rotation, or right rotation, or even both.
Rotations shift the position of the nodes so that o(log N) operations will still be performed, and prevents skewed trees

Cases:

1) Left Left         2) Right Right
O                  O
/                    \
O                      O
/                        \
O                          O

2) Right Left        4) Left Right
O                   O
/                     \
O                       O
\                      /
O                    O

"""
if current_node is None:
return

# Getting the balance factor, which shows if the tree is balanced or not after insertion, by the difference of the heights between the left and right children of the node
balance_factor = self.get_balance_factor(current_node)

# Checks if the node is unbalanced, and if it is, perform one of the 4 cases
# Case 1: Left Left
if balance_factor > 1 and current_node.left:
if value < current_node.left.value:
self.rotate_right(current_node)

# Case 2: Right Right
if balance_factor < -1 and current_node.right:
if value > current_node.right.value:
self.rotate_left(current_node)

# Case 3: Left Right
if balance_factor > 1 and current_node.left:
if value > current_node.left.value:
self.rotate_left(current_node.left)
self.rotate_right(current_node)

# Case 4: Right Left
if balance_factor < -1 and current_node.right:
if value < current_node.right.value:
self.rotate_right(current_node.right)
self.rotate_left(current_node)

self.recursive_balancing_for_insertion(current_node.parent, value)

def recursive_balancing_for_deletion(self, current_node):
"""
As the name says, rebalances, by using balance factor of it's child nodes
It slowly checks the balance factor of every node from the parent of the recently inserted node,
after a node has been checked, it goes the parent of that node, until it reaches the root node, which doesn't have a parent /= None.

Depending on the balance factor and comparison, it may perform a left rotation, or right rotation, or even both.
Rotations shift the position of the nodes so that o(log N) operations will still be performed, and prevents skewed trees

Cases:

1) Left Left         2) Right Right
O                  O
/                    \
O                      O
/                        \
O                          O

2) Right Left        4) Left Right
O                   O
/                     \
O                       O
\                      /
O                    O

"""
if current_node is None:
return

balance_factor = self.get_balance_factor(current_node)

# Case 1 - Left Left
if balance_factor > 1 and self.get_balance_factor(current_node.left) >= 0:
self.rotate_right(current_node)

# Case 2 - Right Right
if balance_factor < -1 and self.get_balance_factor(current_node.right) <= 0:
self.rotate_left(current_node)

# Case 3 - Left Right
if balance_factor > 1 and self.get_balance_factor(current_node.left) < 0:
self.rotate_left(current_node.left)
self.rotate_right(current_node)

# Case 4 - Right Left
if balance_factor < -1 and self.get_balance_factor(current_node.right) > 0:
self.rotate_right(current_node.right)
self.rotate_left(current_node)

self.recursive_balancing_for_deletion(current_node.parent)

def insert(self, value):
"""
Inserts a value into the tree, by traversing through the tree till it finds a None.

1) No balancing is done when inserting the root node.
2) Balancing is done for any other node.
"""
NewNode = Node(value)
current_node = self.root

# Traversing through the Tree till we find the right place to put the value
if current_node is None:
self.root = NewNode
else:
while True:
if value < current_node.value:
NewNode.parent = current_node
#Left
if not current_node.left:
current_node.left = NewNode
break
NewNode.parent = current_node
current_node = current_node.left
else:
NewNode.parent = current_node
#Right
if not current_node.right:
current_node.right = NewNode
break
NewNode.parent = current_node
current_node = current_node.right

# Recursive balancing
self.recursive_balancing_for_insertion(current_node, value)

def lookup(self, value):
"""
Searches for a value in the tree,
if the value is found, returns True
if the value isn't found, returns False

"""
current_node = self.root

if current_node is None:
return None
while current_node:
current_value = current_node.value
if value < current_value:
current_node = current_node.left
elif value > current_value:
current_node = current_node.right
else:
return True
return False

def getMin(self, current_node):
"""
returns the smallest value of a node in the current_node, by going left, until there is no more left child nodes

Used to find the in order successor for a node, incase it has two children when removing.

"""
while current_node.left is not None:
current_node = current_node.left
return current_node.value

def remove(self, value, current_node=None):
"""
Removes a value from the tree
returns True, if value was deleted
returns False, if value wasn't found/deleted

Different position changes depending on 3 cases:
1) Node has no children
easiest to remove, just remove the parent nodes pointer to it
2) Node has 1 child
just point the next child of parent node to the next child of the node to be deleted instead
3) Node has 2 children
hardest to remove, first the in order successor of the node is found, by getting the smallest value in the right child node.
Then the node value is replaced with the in order successor, and the inorder successor is recursively deleted.
"""
if not current_node:
current_node = self.root

if not current_node: # if the root is None
return None

parent_node = None
while current_node:
# value of current node
current_value = current_node.value

# If value to be searched is less than the current node value, then move left
if value < current_value:
parent_node = current_node
current_node = current_node.left

# If value to be searched is greater than the current node value, then move right
elif value > current_value:
parent_node = current_node
current_node = current_node.right

# if value to be searched is the current node's value
else:
# No Child
if not current_node.left and not current_node.right:

# If there is no parent node, then it is the root node
if parent_node is None:
self.root = None
elif current_node == parent_node.left:
parent_node.left = None
else:
parent_node.right = None
self.recursive_balancing_for_deletion(current_node)

# One Child
elif current_node.right and not current_node.left:
if parent_node is None:
self.root = current_node.right
elif current_node == parent_node.left:
parent_node.left = current_node.right
else:
parent_node.right = current_node.right
current_node.right.parent = parent_node
current_node.right = None
self.recursive_balancing_for_deletion(current_node)

elif current_node.left and not current_node.right:
if parent_node is None:
self.root = current_node.left
elif current_node == parent_node.left:
parent_node.left = current_node.left
else:
parent_node.right = current_node.left
current_node.left.parent = parent_node
current_node.left = None
self.recursive_balancing_for_deletion(current_node)

# Two Child
else:
# gets the successor of the current_node
in_order_successor = self.getMin(current_node.right)

# Removes the successor to replace the node to be removed
self.remove(in_order_successor, current_node)

#if the node to be removed is the root node
if parent_node is None:
self.root.value = in_order_successor
elif current_node == parent_node.left:
parent_node.left.value = in_order_successor
else:
parent_node.right.value = in_order_successor

self.recursive_balancing_for_deletion(parent_node)

current_node.parent = None
return True # if removed

return False # if value doesnt exist

test = AVLTree((1, 2, 3, 4, 5))
#test.remove(3)
#test.remove(1)
#test.remove(2)
#test.remove(4)
#test.remove(5)
#print(test.lookup(5))
test.inorder(test.root)


• get_depth() has depth and height backwards. Nov 4, 2021 at 2:13