An AVL Tree implementation

AVLTree.cpp

/**
 * Project/Lab 6
 *
 * AVL Trees
 *
 * Kevin Cuzner
 * CS235
 * Section 4
 */

#include <string>
#include <sstream>

#include <iostream>
using namespace std;

template <typename T>
class AVLTree
{
public:
	AVLTree()
	{
		top = NULL;
		size = 0;
	}

	~AVLTree()
	{
		//traverse the tree and delete the nodes
		if (this->top != NULL)
			delete top;
	}

public:
	class Node
	{
	public:
		Node(T item)
		{
			this->item = item;
			left = NULL;
			right = NULL;
			//since this node has no children by default, both heights are 0
			leftHeight = 0;
			rightHeight = 0;
			changedFlag = true;
		}
		~Node()
		{
			//avoid orphans by deleting the children of this node as well
			if (this->left != NULL)
			{
				delete this->left;
			}
			if (this->right != NULL)
			{
				delete this->right;
			}
		}
		
	protected:
		T item;
		Node* left;
		Node* right;
		bool changedFlag;
		
	public:
		T getItem() { return item; }
		
		int leftHeight;
		int rightHeight;
		
		Node* getRight() { return this->right; }
		Node* getLeft() { return this->left; }
		
		void setLeft(Node* l)
		{
			this->left = l;
			if (this->left == NULL)
			{
				this->leftHeight = 0;
			}
			else
			{
				this->leftHeight = this->left->Height();
			}
			this->changedFlag = true; //our heights have been changed
		}
		
		void setRight(Node* r)
		{
			this->right = r;
			if (this->right == NULL)
			{
				this->rightHeight = 0;
			}
			else
			{
				this->rightHeight = this->right->Height();
			}
			this->changedFlag = true; //our heights have been changed
		}
		
		//replaces a child in this node
		void ReplaceChild(Node* oldChild, Node* newChild)
		{
			if (this->right == oldChild)
			{
				this->setRight(newChild);
			}
			else if (this->left == oldChild)
			{
				this->setLeft(newChild);
			}
			else
			{
				//there was a problem...
				//cout << "Problem replacing child on " << this->item << endl;
			}
		}
		
		//returns the height of this node
		int Height()
		{
			//check to make sure heights are updated
			if (this->right != NULL)
			{
				if (this->right->changedFlag)
				{
					//update our right height
					this->rightHeight = this->right->Height();
					this->right->changedFlag = false; //reset the changed flag
				}
			}
			else
			{
				this->rightHeight = 0;
			}
			
			if (this->left != NULL)
			{
				if (this->left->changedFlag)
				{
					//update our left height
					this->leftHeight = this->left->Height();
					this->left->changedFlag = false; //reset the changed flag
				}
			}
			else
			{
				this->leftHeight = 0;
			}
			
			//get our height
			if (leftHeight > rightHeight)
			{
				return leftHeight + 1;
			}
			else
			{
				return rightHeight + 1;
			}
		}
		
		//returns the balance state of this node
		int Balance()
		{
			if (this->left != NULL)
			{
				leftHeight = this->left->Height();
			}
			else
			{
				leftHeight = 0;
			}
			if (this->right != NULL)
			{
				rightHeight = this->right->Height();
			}
			else
			{
				rightHeight = 0;
			}
			
			//return balance
			return this->leftHeight - this->rightHeight;
		}
	};
	
	//node in a linked list queue used for printing
	struct PrintNode
	{
		PrintNode(PrintNode* prev, PrintNode* nxt, Node* n, int l) :
			previous(prev),
			next(nxt),
			node(n),
			level(l)
		{
		}
		
		//previous node in the list
		PrintNode* previous;
		//next node in the list
		PrintNode* next;
		
		//the node this points to
		Node* node;
		
		//the level on the tree this is
		int level;
	};

public:
	bool Insert(T toInsert, Node* current = NULL, Node* parent = NULL, bool left = false)
	{
		if (parent == NULL)
		{
			//this is the first iteration and so we should set this to the top
			current = this->top;
		}
		
		if (current == NULL)
		{
			//the node gets inserted here
			current = new Node(toInsert);
			if (parent == NULL)
			{
				this->top = current;
			}
			else
			{
				if (left)
					parent->setLeft(current);
				else
					parent->setRight(current);
			}
			this->size++;
			return true; //this iteration is done since our height is just fine
		}
		else
		{
			//continue down the tree
			if (toInsert > current->getItem())
			{
				//it goes to the right
				if (!Insert(toInsert, current->getRight(), current, false))
				{
					//roll the false back up the tree
					return false;
				}
				//update the current item's height on the right
				current->rightHeight = current->getRight()->Height();
			}
			else if (toInsert < current->getItem())
			{
				//it goes to the left
				if (!Insert(toInsert, current->getLeft(), current, true))
				{
					//roll the false back up the tree
					return false;
				}
				//update the current item's height on the left
				current->leftHeight = current->getLeft()->Height();
			}
			else
			{
				//it is already in the tree
				return false; //we are done with this iteration
			}
		}
		
		Balance(current, parent);
		
		return true;
	}

	bool Remove(T toRemove, Node* current = NULL, Node* parent = NULL, bool left = false)
	{
		if (parent == NULL)
		{
			current = this->top;
		}
		
		if (current == NULL)
		{
			//not found
			return false;
		}
		else
		{
			//decide where to go
			if (toRemove < current->getItem())
			{
				//go down the left
				if (!Remove(toRemove, current->getLeft(), current, true))
				{
					return false; //roll the false back up the tree
				}
				//balance the tree
				Balance(current, parent, true);
				return true;
			}
			else if (toRemove > current->getItem())
			{
				//go down the right
				if (!Remove(toRemove, current->getRight(), current, false))
				{
					return false;
				}
				//balance the tree
				Balance(current, parent, true);
				return true;
			}
			else
			{
				//this is the item and we need to remove it
				
				//what is our child status?
				if (current->Height() == 1)
				{
					//we are a leaf
					//replace us with null in our parent
					if (parent == NULL)
					{
						//we were the top...and since we were a leaf now there is no tree
						top = NULL;
					}
					else
					{
						parent->ReplaceChild(current, NULL);
					}
					delete current;
				}
				else if (current->getLeft() == NULL)
				{
					//replace us with our right child
					if (parent == NULL)
					{
						//we were the top...now our right child will be the top
						top = current->getRight();
					}
					else
					{
						parent->ReplaceChild(current, current->getRight());
					}
					current->setRight(NULL); //detach the child from us
					delete current;
				}
				else if (current->getRight() == NULL)
				{
					//repalce us with our left child
					if (parent == NULL)
					{
						//we were the top...now our left child will be the top
						top = current->getLeft();
					}
					else
					{
						parent->ReplaceChild(current, current->getLeft());
					}
					current->setLeft(NULL); //detach the child from us
					delete current;
				}
				else
				{
					//both nodes are full so we will choose the smallest item in the right subtree
					Node* smallest = DetachSmallest(current->getRight(), current);
					
					//smallest will now inherit our children
					smallest->setRight(current->getRight());
					smallest->setLeft(current->getLeft());
					current->setRight(NULL);
					current->setLeft(NULL);
					//and our parent will own it
					if ( parent == NULL )
					{
						//it replaces us as the top
						top = smallest;
						//balance the new node
						Balance(smallest, NULL, true);
					}
					else
					{
						//it replaces us in our parentage
						parent->ReplaceChild(current, smallest);
						//balance the new node
						Balance(smallest, parent, true);
					}
					delete current;
				}
				this->size--;
				return true;
			}
			
			
		}
		
		
		
		return true;
	}

	Node* Find(T item, Node* current = NULL, bool first = true)
	{
		if (first)
		{
			current = top;
		}
		
		if (current == NULL)
		{
			//it wasn't found
			return NULL;
		}
		else if (item < current->getItem())
		{
			//go down the left
			return Find(item, current->getLeft(), false);
		}
		else if (item > current->getItem())
		{
			//go down the right
			return Find(item, current->getRight(), false);
		}
		else
		{
			//we found it!
			return current;
		}
	}
	
	//clears the tree of everything
	void Clear()
	{
		//quickly clean the tree
		if (this->top != NULL)
			delete top;
		top = NULL;
		//well that was easy
	}

	//prints out the tree into a string
	std::string Print()
	{
		if (this->top == NULL)
		{
			return ""; //nothing to print
		}
		
		stringstream sb;
		
		//create a queue-like structure
		PrintNode* front = new PrintNode(NULL, NULL, top, 0);
		PrintNode* back = front;
		
		//iterate through
		int numOnLevel = 0;
		int level = 0;
		sb << "Level 0:";
		while(front != NULL)
		{
			Node* current = front->node;
			
			if (front->level != level)
			{
				//we are in a new level
				sb << endl << "Level " << front->level << ":";
				level = front->level;
				numOnLevel = 0;
			}
			
			if (numOnLevel == 8)
			{
				//there are already 8 items on the level. Re-output the level thing and reset the counter
				sb << endl << "Level " << front->level << ":";
				numOnLevel = 0;
			}
			
			//de-queue the front
			PrintNode* newFront = front->next;
			if (newFront != NULL)
			{
				//this no longer has a previous node
				newFront->previous = NULL;
			}
			delete front;
			front = newFront; //the front is now this node
			
			//print out this node?
			sb << " " << current->getItem() << " (" << current->Height() << ")";
			numOnLevel++;
			
			//enqueue its children to the back if needed
			if (current->getLeft() != NULL)
			{
				if (front == NULL)
				{
					//the front is the back right now
					front = back = new PrintNode(NULL, NULL, current->getLeft(), level + 1);
				}
				else
				{
					//the back will move along as normal
					back->next = new PrintNode(back, NULL, current->getLeft(), level + 1);
					back = back->next;
				}
			}
			if (current->getRight() != NULL)
			{
				if (front == NULL)
				{
					//the front is the back right now
					front = back = new PrintNode(NULL, NULL, current->getRight(), level + 1);
				}
				else
				{
					//the back will move along as normal
					back->next = new PrintNode(back, NULL, current->getRight(), level + 1);
					back = back->next;
				}
			}
		}
		
		sb << endl;
		
		return sb.str();
	}

public:
	Node* top;
	int size;
	
	void Balance(Node* current, Node* parent, bool removal = false)
	{
		//cout << "Balancing " << current->getItem() << " " << current->Balance() << endl;
		//check for a rotational need
		int balance = current->Balance();
		if (balance < -1)
		{
			//it is heavy on the right side
			//cout << current->getItem() << " is heavy on the right" << endl;
			balance = current->getRight()->Balance();
			//cout << balance << endl;
			if (balance < 0)
			{
				//single left rotation is needed with the current node as the root
				Node* n = RotateLeft(current);
				if (parent != NULL)
				{
					//re-attach this node to what used to be the parent of the current node
					parent->ReplaceChild(current, n);
				}
				else
				{
					//this node is the new top
					this->top = n;
				}
			}
			else if (balance > 0)
			{
				//rotate right with current->right as the root and then rotate left with current as the root
				Node* n = RotateRight(current->getRight());
				current->ReplaceChild(current->getRight(), n);
				n = RotateLeft(current);
				if (parent != NULL)
				{
					//re-attach this node to what used to be the parent of current
					parent->ReplaceChild(current, n);
				}
				else
				{
					//this node is the new top
					this->top = n;
				}
			}
			else if (removal)
			{
				//we are removing something that has equal height grandchildren.
				//current must become the left child of current->right and current->right's left child must become the right child of current
				//cout << "Equal height grandchildren" << endl;
				if (parent != NULL)
				{
					parent->ReplaceChild(current, current->getRight());
				}
				else
				{
					//replacing the top
					top = current->getRight();
				}
				Node* temp = current->getRight();
				current->setRight(current->getRight()->getLeft());
				temp->setLeft(current);
			}
		}
		else if (balance > 1)
		{
			//cout << current->getItem() << " is heavy on the left" << endl;
			//it is heavy on the left side
			balance = current->getLeft()->Balance();
			//cout << balance << endl;
			if (balance > 0)
			{
				//single right rotation is needed with the current node as the root
				Node* n = RotateRight(current);
				if (parent != NULL)
				{
					//re-attach this node to what used to be the parent of the current node
					parent->ReplaceChild(current, n);
				}
				else
				{
					//this node is the new top
					this->top = n;
				}
			}
			else if (balance < 0)
			{
				//rotate left with current->right as the root and then rotate right with current as the root
				Node* n = RotateLeft(current->getLeft());
				current->ReplaceChild(current->getLeft(), n);
				n = RotateRight(current);
				if (parent != NULL)
				{
					//re-attach this node to what used to be the parent of current
					parent->ReplaceChild(current, n);
				}
				else
				{
					//this node is the new top
					this->top = n;
				}
			}
			else if (removal)
			{
				//we are removing something that has equal height grandchildren.
				//cout << "Equal height grandchildren" << endl;
				if (parent != NULL)
				{
					parent->ReplaceChild(current, current->getLeft());
				}
				else
				{
					//replacing the top
					top = current->getLeft();
				}
				Node* temp = current->getLeft();
				current->setLeft(current->getLeft()->getRight());
				temp->setRight(current);
			}
		}
	}

	Node* RotateLeft(Node* n)
	{
		Node* k = n->getRight();
		
		n->setRight(k->getLeft());
		k->setLeft(n);
		
		return k;
	}

	Node* RotateRight(Node* n)
	{
		Node* k = n->getLeft();
		
		n->setLeft(k->getRight());
		k->setRight(n);
		
		return k;
	}
	
	//this will detach the smallest node in a tree and return it
	Node* DetachSmallest(Node* current, Node* parent)
	{
		if (parent == NULL)
		{
			//this shouldn't ever happen...there was an error
			return NULL;
		}
		
		//cout << "Looking at " << parent->getItem() << "'s child" << endl;
		
		if (current == NULL)
		{
			return NULL;
		}
		
		//attempt to traverse our left subtree
		Node* detached = DetachSmallest(current->getLeft(), current);
		if (detached == NULL)
		{
			//we are the node to be detached! Our right child will replace us
			parent->ReplaceChild(current, current->getRight());
			if (current->getRight() != NULL)
			{
				Balance(current->getRight(), parent, true);
			}
			current->setRight(NULL); //we no longer have a right child
			return current;
		}
		else
		{
			//someone somewhere was detached, so let's roll them up the line
			//balance the tree here?
			//cout << current->getItem() << "'s left is " << current->leftHeight << " and their right is " << current->rightHeight << endl;
			Balance(current, parent, true);
			return detached;
		}
	}
};