AVL Tree Implementation in Python
An AVL tree is a self-balancing binary search tree where the difference between heights of left and right subtrees (the balance factor) is at most one for all nodes. If at any time during insertion or deletion, this condition is not met, the tree is rebalanced using rotations. The AVL tree provides efficient search, insertion, and deletion operations, all of which run in O(log n) time.
Program Structure
The AVL tree implementation in Python consists of the following key components:
1. Node Class
This class represents a node in the AVL tree. It includes the following attributes:
key: The key stored in the node.height: The height of the node in the tree.left,right: Pointers to the left and right child nodes, respectively.
2. AVLTree Class
This class encapsulates the functionality of the AVL tree and includes the following methods:
insert(self, root, key): Inserts a new key into the AVL tree, rebalancing the tree if necessary.delete(self, root, key): Deletes a key from the AVL tree, rebalancing the tree if necessary.left_rotate(self, z): Performs a left rotation around a given node.right_rotate(self, z): Performs a right rotation around a given node.get_height(self, node): Returns the height of a given node.get_balance(self, node): Returns the balance factor of a given node.min_value_node(self, node): Finds the node with the minimum key in a given subtree.pre_order(self, root): Performs a pre-order traversal of the tree.
Python Code Implementation
class Node:
def __init__(self, key):
"""
Initialize a new node.
:param key: The key value of the node.
"""
self.key = key
self.left = None
self.right = None
self.height = 1 # New node is initially added at leaf, height = 1
class AVLTree:
def get_height(self, node):
"""
Get the height of the node.
:param node: The node to get the height of.
:return: Height of the node, or 0 if the node is None.
"""
if not node:
return 0
return node.height
def get_balance(self, node):
"""
Get the balance factor of the node.
:param node: The node to get the balance factor of.
:return: Balance factor of the node, or 0 if the node is None.
"""
if not node:
return 0
return self.get_height(node.left) - self.get_height(node.right)
def left_rotate(self, z):
"""
Perform a left rotation around the node z.
:param z: The node to rotate around.
:return: New root after rotation.
"""
y = z.right
T2 = y.left
# Perform rotation
y.left = z
z.right = T2
# Update heights
z.height = 1 + max(self.get_height(z.left), self.get_height(z.right))
y.height = 1 + max(self.get_height(y.left), self.get_height(y.right))
# Return the new root
return y
def right_rotate(self, z):
"""
Perform a right rotation around the node z.
:param z: The node to rotate around.
:return: New root after rotation.
"""
y = z.left
T3 = y.right
# Perform rotation
y.right = z
z.left = T3
# Update heights
z.height = 1 + max(self.get_height(z.left), self.get_height(z.right))
y.height = 1 + max(self.get_height(y.left), self.get_height(y.right))
# Return the new root
return y
def insert(self, root, key):
"""
Insert a key into the AVL tree.
:param root: The root node of the AVL tree.
:param key: The key to insert.
:return: The new root of the subtree after insertion.
"""
# 1. Perform the normal BST insertion
if not root:
return Node(key)
elif key < root.key:
root.left = self.insert(root.left, key)
else:
root.right = self.insert(root.right, key)
# 2. Update height of this ancestor node
root.height = 1 + max(self.get_height(root.left), self.get_height(root.right))
# 3. Get the balance factor
balance = self.get_balance(root)
# 4. If the node becomes unbalanced, then balance it
# Left Left Case
if balance > 1 and key < root.left.key:
return self.right_rotate(root)
# Right Right Case
if balance < -1 and key > root.right.key:
return self.left_rotate(root)
# Left Right Case
if balance > 1 and key > root.left.key:
root.left = self.left_rotate(root.left)
return self.right_rotate(root)
# Right Left Case
if balance < -1 and key < root.right.key:
root.right = self.right_rotate(root.right)
return self.left_rotate(root)
return root
def min_value_node(self, node):
"""
Get the node with the minimum key value found in that tree.
:param node: The node to start the search from.
:return: Node with the minimum key.
"""
current = node
while current.left is not None:
current = current.left
return current
def delete(self, root, key):
"""
Delete a key from the AVL tree.
:param root: The root node of the AVL tree.
:param key: The key to delete.
:return: The new root of the subtree after deletion.
"""
# 1. Perform standard BST delete
if not root:
return root
elif key < root.key:
root.left = self.delete(root.left, key)
elif key > root.key:
root.right = self.delete(root.right, key)
else:
# Node with only one child or no child
if root.left is None:
temp = root.right
root = None
return temp
elif root.right is None:
temp = root.left
root = None
return temp
# Node with two children: Get the inorder successor
temp = self.min_value_node(root.right)
# Copy the inorder successor's content to this node
root.key = temp.key
# Delete the inorder successor
root.right = self.delete(root.right, temp.key)
# If the tree had only one node, return it
if root is None:
return root
# 2. Update the height of the current node
root.height = 1 + max(self.get_height(root.left), self.get_height(root.right))
# 3. Get the balance factor
balance = self.get_balance(root)
# 4. If the node becomes unbalanced, then balance it
# Left Left Case
if balance > 1 and self.get_balance(root.left) >= 0:
return self.right_rotate(root)
# Left Right Case
if balance > 1 and self.get_balance(root.left) < 0:
root.left = self.left_rotate(root.left)
return self.right_rotate(root)
# Right Right Case
if balance < -1 and self.get_balance(root.right) <= 0:
return self.left_rotate(root)
# Right Left Case
if balance < -1 and self.get_balance(root.right) > 0:
root.right = self.right_rotate(root.right)
return self.left_rotate(root)
return root
def pre_order(self, root):
"""
Perform a pre-order traversal of the AVL tree.
:param root: The root of the AVL tree.
"""
if not root:
return
print(f"{root.key} ", end="")
self.pre_order(root.left)
self.pre_order(root.right)
Explanation
The AVL tree implementation revolves around two main classes: Node and AVLTree.
1. Node Class
The Node class represents a node in the AVL tree, containing:
key: The key stored in the node.height: The height of the node, which is crucial for maintaining balance in the AVL tree.left,right: Pointers to the left and right child nodes, respectively.
2. AVLTree Class
The AVLTree class handles the operations of the AVL tree, including insertion, deletion, and tree balancing.
Insertion
The insert method adds a new key to the AVL tree. After insertion, the method checks if the tree is balanced. If any node becomes unbalanced, the appropriate rotations (left, right, left-right, or right-left) are performed to restore balance.
Deletion
The delete method removes a key from the AVL tree. Similar to insertion, after deletion, the tree is checked for balance, and rotations are applied if necessary.
Rotations
The left_rotate and right_rotate methods perform rotations to adjust the tree structure. Rotations are necessary to maintain the balanced nature of the AVL tree during insertion and deletion operations.
Balance Factor
The get_balance method calculates the balance factor of a node, which is the difference in height between the left and right subtrees. The balance factor is used to determine if a node is balanced.
Traversal
The pre_order method allows a pre-order traversal of the AVL tree, which is useful for checking the structure of the tree.
Usage Example
# Example usage of the AVLTree
# Create an AVL tree
avl_tree = AVLTree()
# Insert some keys into the tree
root = None
keys = [10, 20, 30, 40, 50, 25]
for key in keys:
root = avl_tree.insert(root, key)
# Perform pre-order traversal
print("Pre-order traversal of the constructed AVL tree is:")
avl_tree.pre_order(root)
# Output should show the keys in pre-order traversal
This example demonstrates how to create an AVL tree, insert keys into it, and perform a pre-order traversal to inspect the structure of the tree.
