I recently was asked to do an AVL Tree implementation in C/C++. I've kept this implementation rather simple with no AVL Tree class and focus on the algorithm itself. So the tree is represented by a pointer to the HEAD Node and modifying functions take a double pointer to this node.
The main() function tests the tree by adding and deleting random numbers into and from the tree.
/*
* Tree.cpp
* AVL Tree - Always balanced with maximum left and right sub-tree height difference of 1!
* Created on: Oct 14, 2011
* Author: Fahd
*/
#include <iostream>
using namespace std;
// utility
#define MAX(X,Y) (X>Y?X:Y)
// node structure
struct node {
int data;
node* left;
node* right;
int height; // interim use .. value not necessarily consistent with state
};
// returns pointer to pointer to node with "data", if not then NULL
node** avlLookup(node** tree, int data) {
node** current = tree;
while (*current) {
if ((*current)->data == data) return current;
if (data < (*current)->data) current = &((*current)->left);
else current = &((*current)->right);
}
return NULL;
}
// new node with required data set
node* newNode(int data) {
node* n = new node;
n->data = data;
n->left = NULL;
n->right = NULL;
n->height = -1;
return n;
}
// computes tree heights at each node (given subtree heights are done)
int avlComputeHeight(node* tree) {
if (!tree) return 0;
int leftHeight = (tree->left)?(tree->left->height):0;
int rightHeight = (tree->right)?(tree->right->height):0;
tree->height = MAX(leftHeight,rightHeight)+1;
return tree->height;
}
// returns difference between left and right subtree heights
int avlHeightDifference(node** ppNode) {
if (!(*ppNode)) return 0;
int rightHeight = ((*ppNode)->right)?(*ppNode)->right->height:0;
int leftHeight = ((*ppNode)->left)?(*ppNode)->left->height:0;
return rightHeight - leftHeight;
}
// Balances an unbalanced tree
void avlRebalance(node** balancePoint) {
if (!(*balancePoint)) return;
// balance subtrees first
avlRebalance(&((*balancePoint)->right));
avlRebalance(&((*balancePoint)->left));
// compute height of self node
avlComputeHeight(*balancePoint);
// if height difference is less than 2 then return
if (abs(avlHeightDifference(balancePoint))<2)
return;
// test which subtree is heavier
if (avlHeightDifference(balancePoint)>0) {
// if right subtree is heavier
node* newRoot = (*balancePoint)->right;
node* newLeftRight = (*balancePoint)->right->left;
node* newLeft = *balancePoint;
*balancePoint = newRoot;
newRoot->left = newLeft;
newRoot->left->right = newLeftRight;
} else {
// if left subtree is heavier
node* newRoot = (*balancePoint)->left;
node* newRightLeft = (*balancePoint)->left->right;
node* newRight = *balancePoint;
*balancePoint = newRoot;
newRoot->right = newRight;
newRoot->right->left = newRightLeft;
}
}
// inserts a new node into tree and balances if needed
void avlInsert(node** tree, int data) {
node** insertionPoint = tree;
// find where to insert into tree
while (*insertionPoint) {
if (data == (*insertionPoint)->data) {
// if already in tree, ignore insert
cout<<"Data ["<<data<<"] already in tree, ignoring insert"<<endl;;
return;
}
if (data < (*insertionPoint)->data) {
// insert into left subtree
insertionPoint = &((*insertionPoint)->left);
}
else {
// insert into right subtree
insertionPoint = &((*insertionPoint)->right);
}
}
// insert new node
*insertionPoint = newNode(data);
// balance if needed
avlRebalance(tree);
return;
}
// recursive print tree (for testing only)
void avlPrint(node* tree) {
if (!tree) return;
cout<<"Node: "<<tree->data;
cout<<", Left: ";
if (tree->left) cout<<tree->left->data;
else cout<<"-";
cout<<", Right: ";
if (tree->right) cout<<tree->right->data;
else cout<<"-";
cout<<endl;
avlPrint(tree->left);
avlPrint(tree->right);
}
// get right most node (predecessor of node about to be deleted)
node** avlRightMost(node** tree) {
node** result = tree;
while ((*result)->right) {
result = &((*result)->right);
}
return result;
}
// delete node with "data" in "tree"
void avlDelete(node** tree, int data) {
if (!(*tree)) return;
// find which node to delete
node** ppRequiredNode = avlLookup(tree, data);
// if node found
if (ppRequiredNode) {
node* pNode = *ppRequiredNode;
if (!(pNode->left) && !(pNode->right)) {
// if leaf node
delete pNode;
*ppRequiredNode = NULL;
} else if (pNode->left && pNode->right) {
// if node has both left and right subtrees
node** rightMost = avlRightMost(&(pNode->left));
pNode->data = (*rightMost)->data;
node* toDel = *rightMost;
*rightMost = (*rightMost)->left;
delete toDel;
} else if (pNode->left) {
// if only left subtree exists
*ppRequiredNode = pNode->left;
delete pNode;
} else {
// if only right subtree exists
*ppRequiredNode = pNode->right;
delete pNode;
}
// rebalance if needed after delete
avlRebalance(tree);
} else {
// if node is not in tree
cout<<"Nothing to delete."<<endl;
}
}
// main function to setup test environment
int main() {
// make tree
node* tree = NULL;
srand(0);
// randomly insert and delete nodes into and from tree (nodes 0~14)
for (int i=0; i<100; i++) {
if (rand()%2) {
avlInsert(&tree, rand()%15);
} else {
avlDelete(&tree, rand()%15);
}
}
// print tree after operations
avlPrint(tree);
// lookup all nodes in tree
for (int reqData = 0; reqData<15; reqData++) {
node** requiredNode = avlLookup(&tree, reqData);
if (requiredNode) {
cout<<reqData<<" Found."<<endl;
} else {
cout<<reqData<<" Not found."<<endl;
}
}
return 0;
}
