#include<iostream>
using namespace std;
template <typename T>
class AVL {
public:
// Node class represents a node in the AVL tree
class node {
public:
T key; // Data of the node
int height; // Height of the node
node* left; // Pointer to the left child
node* right; // Pointer to the right child
// Constructor to initialize the node
node(T k) : key(k), height(1), left(NULL), right(NULL) {} // height is initialized to 1 for leaf nodes
};
node* root = NULL; // Root of the AVL tree
// Public methods for the AVL tree
void insert(T x) {
root = insertUtil(root, x); // Insert value into the tree
}
void remove(T x) {
root = removeUtil(root, x); // Remove value from the tree
}
node* search(T x) {
return searchUtil(root, x); // Search for a value in the tree
}
void inorder() {
inorderUtil(root); // Perform inorder traversal
cout << endl;
}
private:
// Private helper methods for AVL tree operations
// Returns the height of a given node
int height(node* head) {
if (head == NULL)
return 0;
return head->height;
}
// Updates the height of a node based on its children
void updateHeight(node* head) {
head->height = 1 + max(height(head->left), height(head->right));
}
// Calculates the balance factor of a node (left subtree height - right subtree height)
int balanceFactor(node* head) {
if (head == NULL)
return 0;
return (height(head->left) - height(head->right));
}
// Right rotation to balance the tree (used when left-heavy)
node* rightRotation(node* head) {
node* newhead = head->left;
head->left = newhead->right;
newhead->right = head;
updateHeight(head); // Update the height of the old root
updateHeight(newhead); // Update the height of the new root
return newhead;
}
// Left rotation to balance the tree (used when right-heavy)
node* leftRotation(node* head) {
node* newhead = head->right;
head->right = newhead->left;
newhead->left = head;
updateHeight(head); // Update the height of the old root
updateHeight(newhead); // Update the height of the new root
return newhead;
}
// Helper function to insert a value into the AVL tree
node* insertUtil(node* head, T x) {
if (!head) return new node(x); // Create a new node if the tree is empty
// Recursively insert the value into the correct position in the tree
if (x < head->key)
head->left = insertUtil(head->left, x);
else if (x > head->key)
head->right = insertUtil(head->right, x);
// Update the height of the current node
updateHeight(head);
// Check balance factor and perform rotations if necessary
int balance = balanceFactor(head);
// Left-Left case
if (balance > 1 && x < head->left->key)
return rightRotation(head);
// Right-Right case
if (balance < -1 && x > head->right->key)
return leftRotation(head);
// Left-Right case
if (balance > 1 && x > head->left->key) {
head->left = leftRotation(head->left);
return rightRotation(head);
}
// Right-Left case
if (balance < -1 && x < head->right->key) {
head->right = rightRotation(head->right);
return leftRotation(head);
}
return head; // Return the (possibly rotated) node
}
// Helper function to remove a value from the AVL tree
node* removeUtil(node* head, T x) {
if (!head)
return NULL; // Base case: value not found
if (x < head->key) {
head->left = removeUtil(head->left, x);
} else if (x > head->key) {
head->right = removeUtil(head->right, x);
} else {
// Node to be deleted is found
if (!head->left || !head->right) {
node* temp = head->left ? head->left : head->right;
delete head;
head = temp;
} else {
// Node has two children
node* temp = minNode(head->right);
head->key = temp->key;
head->right = removeUtil(head->right, temp->key);
}
}
if (!head) return NULL; // Return NULL if node was deleted
// Update the height of the current node
updateHeight(head);
// Rebalance the tree if needed
int balance = balanceFactor(head);
// Left-Left case
if (balance > 1) {
if (balanceFactor(head->left) >= 0)
return rightRotation(head);
else {
head->left = leftRotation(head->left);
return rightRotation(head);
}
}
// Right-Right case
if (balance < -1) {
if (balanceFactor(head->right) <= 0)
return leftRotation(head);
else {
head->right = rightRotation(head->right);
return leftRotation(head);
}
}
return head; // Return the (possibly rotated) node
}
// Helper function to search for a value in the AVL tree
node* searchUtil(node* head, T x) {
if (!head || head->key == x)
return head; // Base case: value found or tree is empty
if (x < head->key)
return searchUtil(head->left, x);
return searchUtil(head->right, x);
}
// Helper function to perform inorder traversal and print the tree
void inorderUtil(node* head) {
if (head) {
inorderUtil(head->left); // Traverse the left subtree
cout << head->key << " "; // Print the key of the current node
inorderUtil(head->right); // Traverse the right subtree
}
}
// Find the minimum value node in the tree
node* minNode(node* head) {
while (head && head->left)
head = head->left;
return head;
}
};
int main() {
AVL<int> tree; // Create an AVL tree object
int choice, value;
do {
// Display menu options
cout << "Menu:" << endl;
cout << "1. Insert" << endl;
cout << "2. Remove" << endl;
cout << "3. Search" << endl;
cout << "4. Inorder Traversal" << endl;
cout << "5. Exit" << endl;
cout << "Enter your choice: ";
cin >> choice;
switch(choice) {
case 1:
cout << "Enter value to insert: ";
cin >> value;
tree.insert(value); // Insert value into the tree
cout << "Value inserted successfully!" << endl;
break;
case 2:
cout << "Enter value to remove: ";
cin >> value;
tree.remove(value); // Remove value from the tree
cout << "Value removed successfully!" << endl;
break;
case 3:
cout << "Enter value to search: ";
cin >> value;
if (tree.search(value)) { // Search for a value
cout << "Value found!" << endl;
} else {
cout << "Value not found!" << endl;
}
break;
case 4:
cout << "Inorder Traversal: ";
tree.inorder(); // Print the tree using inorder traversal
break;
case 5:
cout << "Exiting program..." << endl;
break;
default:
cout << "Invalid choice, please try again." << endl;
}
} while(choice != 5); // Loop until the user chooses to exit
return 0;
}
