C++ 自平衡二叉搜索树的实现

头文件

// AVLTree.h  for class avltree

#ifndef AVLTREE_H__________
#define AVLTREE_H__________
//#include <AFX.H>

typedef int BOOL;
#define NULL 0
#define TRUE 1
#define FALSE 0
typedef int ElemType;
struct TreeNode    // AVL树数据结构
{
	ElemType data;
	int color;
	struct TreeNode *parent;
	struct TreeNode *left;
	struct TreeNode *right;

	TreeNode(ElemType val);
	int GetLeftHight();
	int GetRightHight();
	int GetDiffHight();
	void GetNodes(TreeNode *pNode, int h);
};

// 栈类 用于存放查找结点的路径，便于回溯找到最小不平衡结点
class CStack
{
public:
	CStack(int nMaxSize = 100);
	~CStack();

	BOOL Pop(TreeNode *& node);
	BOOL Push(TreeNode * node);
	BOOL Top(TreeNode *&node);
	BOOL GetNext(TreeNode* &node);
	int Size()const{return m_nCount; }

private:
	int m_nCount;
	int m_nCurPos;
	int m_nMaxSize;
	TreeNode* *m_arrNode;  // 存放数据结点
};

// AVL树类 自平衡二叉搜索树
class CAVLTree
{
public:
	CAVLTree();
	~CAVLTree();

	int Insert(ElemType val);
	int Remove(ElemType val);
	TreeNode* Find(const ElemType &val);
	void DoBalance(TreeNode *& pCurNode, TreeNode *& pParent, int nLRFlag);
	void PrintTree();
	void EmptyTree();
	int GetSize()const { return m_nSize; }

private:
	void GetNoBalanceNode(CStack &st, TreeNode *& pCurNode, TreeNode *& pParent, int &nLRFlag);
	BOOL LL_Rotate(TreeNode *& pCurNode, TreeNode *& pParentNode, const int &nLRFlag);
	BOOL LR_Rotate(TreeNode *& pCurNode, TreeNode *& pParentNode, const int &nLRFlag);
	BOOL RR_Rotate(TreeNode *& pCurNode, TreeNode *& pParentNode, const int &nLRFlag);
	BOOL RL_Rotate(TreeNode *& pCurNode, TreeNode *& pParentNode, const int &nLRFlag);
	void PrintTree(TreeNode *& pNode);
	void GetNodes(CStack &st, TreeNode *pNode);
	void SwapNodes(TreeNode *&pNode1, TreeNode *&pNode2);

private:
	TreeNode *m_root;
	int m_nSize;
};

#endif

实现文件

// AVLTree.cpp  for implete class CAVLTree
#include <iostream>
#include "AVLTree.h"
#include <Windows.h>
using namespace std;
///////////////////////

/// constructure functions

//////////////////////////////////////////////// struct TreeNode implete
//////////////
TreeNode::TreeNode( ElemType val )  // 构造函数
{
	data = val;
	left = NULL;
	right = NULL;
	parent = NULL;
	color = 0;
}
int nHight = 0;

int TreeNode::GetLeftHight()
{ // 获取结点左子树高度
	nHight = 0;
	GetNodes(this->left, 0);
	return nHight;
}

int TreeNode::GetRightHight()
{  // 获取结点右子树高度
	nHight = 0;
	GetNodes(this->right, 0);
	return nHight;
}

int TreeNode::GetDiffHight()
{  // 获取结点左右子树高度的差，当放回值的绝对值大于1时，该结点为不平衡结点，需要进行结点的自平衡处理
	return (GetLeftHight() - GetRightHight());
}

void TreeNode::GetNodes(TreeNode *pNode, int h)
{ // 递归直到搜索到叶子结点，取最大的值为某结点到叶子结点的高度
	if (pNode == NULL)
	{
		nHight = (nHight > h)? nHight:h;
		return;
	}
	h++;
	GetNodes(pNode->left, h);
	GetNodes(pNode->right, h);
}

///////////////////////////////////////////////////////// class CStack implete
/////////////////

CStack::CStack( int nMaxSize )  // 构造函数
{
	m_nCount = 0;
	m_nMaxSize = nMaxSize;
	m_arrNode = new TreeNode *[m_nMaxSize];
}

CStack::~CStack()  // 析构函数
{
	delete []m_arrNode;
}

BOOL CStack::Pop( TreeNode *& node )  // 出栈
{
	if (m_nCount <= 0)
		return FALSE;
	m_nCount--;
	node = m_arrNode[m_nCount];
	return TRUE;
}

BOOL CStack::Push( TreeNode * node )  // 进栈
{
	if(m_nCount < m_nMaxSize-1)
	{
		m_arrNode[m_nCount] = node;
		m_nCount++;
		return TRUE;
	}
	return FALSE;
}

BOOL CStack::Top( TreeNode *&node )  // 获取栈顶值
{
	if (m_nCount > 0)
	{
		node = m_arrNode[m_nCount-1];
		m_nCurPos = m_nCount - 1;
		return TRUE;
	}
	return FALSE;
}

BOOL CStack::GetNext( TreeNode* &node ) // 第一次返回的值为栈顶的下一数值
{
	if (m_nCurPos-1 < m_nCount-1 && m_nCurPos-1 >= 0)
	{
		m_nCurPos--;
		node = m_arrNode[m_nCurPos];
		return TRUE;
	}
	return FALSE;
}


///////////////////////////////////////////////////////// class CAVLTree implete
//////////////////////
CAVLTree::CAVLTree()  // 构造函数
{
	m_root = NULL;
	m_nSize = 0;
}

CAVLTree::~CAVLTree()
{
	EmptyTree();  //释放该树结构所有结点
}

int CAVLTree::Insert( ElemType val )  // 插入树结点
{
	CStack stNodes(100000);
	TreeNode *pNewNode = new TreeNode(val);
	TreeNode *pNode = m_root;
	TreeNode *preNode = NULL;
	int nFlag;  // 0--left 1--right
	if (m_root == NULL)
	{  // 该树为空
		m_root = pNewNode;
		cout<< "Insert root node: " << pNewNode->data <<" succeed." <<endl;
	//	Sleep(500);
		m_nSize++;
		return 1;
	}
	while(1)
	{   // 查找要插入的位置
		preNode = pNode;
		if (val < pNode->data)
		{
			stNodes.Push(pNode);
			pNode = pNode->left;
			nFlag = 0;
		}
		else if (val > pNode->data)
		{
			stNodes.Push(pNode);
			pNode = pNode->right;
			nFlag = 1;
		}
		else 
		{
			delete pNewNode;
			return 0;
		}

		if (!pNode) //页子结点
		{
			if (nFlag == 0)
			{
				preNode->left = pNewNode;
			}
			else
				preNode->right = pNewNode;
	//		Sleep(500);
			break;
		}
	}
	TreeNode *pNBLNode = NULL;
	TreeNode *pNBLParent = NULL;
	nFlag = 0;
	GetNoBalanceNode(stNodes, pNBLNode, pNBLParent, nFlag);
	if (pNBLNode)
	{
		DoBalance(pNBLNode, pNBLParent, nFlag);
	}
	m_nSize++;
	cout<< "Insert "<< m_nSize <<" node: " << pNewNode->data <<" succeed." <<endl; 
	return 1;
}

int CAVLTree::Remove( ElemType val )  // 删除某一结点
{
	if (!m_root) return 0;
	CStack stNodes(100000);  // 记录查找的路径结点
	TreeNode *pNode = m_root;
	TreeNode *pPreNode = m_root;
	int nLRFlag;
	while(pNode)
	{
		if (val < pNode->data)
		{
			stNodes.Push(pNode);
			if (pNode->left)
			{
				pPreNode = pNode;
			}
			pNode = pNode->left;
			nLRFlag = 0;
		}
		else if (val > pNode->data)
		{
			stNodes.Push(pNode);
			if (pNode->right)
			{
				pPreNode = pNode;
			}
			pNode = pNode->right;
			nLRFlag = 1;
		}
		else break;	
	}
	if (!pNode)   //没找到要删除的结点
		return 0;
	if (!stNodes.Top(pPreNode))
	{ // tree has only one root node.
		pPreNode = NULL;
		if (pNode->left == NULL && pNode->right == NULL)
		{
			delete pNode;
			m_root = NULL;
			m_nSize--;
			return 1;
		}
	}
	if (pNode->left)  //pNode为要删除的结点
	{ // 寻找该结点的左子树中最大的结点与之交换位置
		TreeNode *pNodeDel = pNode->left;
		TreeNode *pPre = pNode;
		stNodes.Push(pPre);
		while(pNodeDel->right)
		{
			pPre = pNodeDel;
			stNodes.Push(pPre);
			pNodeDel = pNodeDel->right;
		}
		SwapNodes(pNode, pNodeDel);  // 将要删除的结点与符合要求的叶子结点交换，然后删除该叶子结点
		if (pPre->left == pNodeDel)// pNode 左边只有一个叶子结点
		{
			pPre->left = pNodeDel->left;
		}
		else
			pPre->right = pNodeDel->left;

		delete pNodeDel;  // 释放叶子结点
	}
	else if (pNode->right)
	{ // 寻找该结点的右子树中最小的结点与之交换位置
		TreeNode *pNodeDel = pNode->right;
		TreeNode *pPre = pNode;
		stNodes.Push(pPre);
		while(pNodeDel->left)
		{
			pPre = pNodeDel;
			stNodes.Push(pPre);
			pNodeDel = pNodeDel->left;
		}
		SwapNodes(pNode, pNodeDel);
		pPre->right = pNodeDel->left;

		delete pNodeDel;
	}
	else  // 该结点为叶子结点
	{
		delete pNode;
		if (nLRFlag == 0)
		{
			pPreNode->left = NULL;
		}
		else
			pPreNode->right = NULL;
	}
	cout << "Remove node: " << val << " succeeded!\n";
//	Sleep(500);
	/// 检查不平衡结点
	TreeNode *pNBLNode = NULL;
	TreeNode *pPreNBLNode = NULL;
	nLRFlag = 0;
	GetNoBalanceNode(stNodes, pNBLNode, pPreNBLNode, nLRFlag);
	if (pNBLNode)
	{
		DoBalance(pNBLNode, pPreNBLNode, nLRFlag);
	}
	m_nSize--;
	return 1;	
}

TreeNode* CAVLTree::Find( const ElemType &val )  // 查找某一结点
{
	if (m_root == NULL) return NULL;
	TreeNode *pNode = m_root;
	while (pNode)
	{
		if (pNode->data < val)
		{
			pNode = pNode->right;
		}
		else if (pNode->data > val)
		{
			pNode = pNode->left;
		}
		else
		{
			cout << "Find node: " << val << " succeeded!\n";
			return pNode;
		}
	}
	return NULL;	
}

// 获取不平衡的结点的信息
void CAVLTree::GetNoBalanceNode( CStack &st, TreeNode *& pCurNode, TreeNode *& pParent, int &nLRFlag )
{
	TreeNode *pNode;
	if (!st.Top(pNode))	return ;
	while(1)
	{
		int nDiff = pNode->GetDiffHight();
		if (nDiff > 1 || nDiff < -1)
		{
			pCurNode = pNode;
			if (st.GetNext(pParent))
			{
				if (pParent->left == pCurNode)
				{
					nLRFlag = 0;  // left child
				}
				else /*if (pParent->right == pCurNode)*/
				{
					nLRFlag = 1;  // right child
				}
				return ;
			}
		}
		if (!st.GetNext(pNode))
		{
			return ;
		}
	}
}

// 自平衡处理函数
void CAVLTree::DoBalance( TreeNode *& pCurNode, TreeNode *& pParent, int nLRFlag )
{
	cout << "Call DoBalance() succeeded!\n";
	if (pCurNode->GetLeftHight() > pCurNode->GetRightHight())
	{
		if (pCurNode->left->GetLeftHight() > pCurNode->left->GetRightHight())
		{
			LL_Rotate(pCurNode, pParent, nLRFlag);  // LR型 平衡处理
		}
		else
			LR_Rotate(pCurNode, pParent, nLRFlag);  // LL型 平衡处理
	}
	else
	{
		if (pCurNode->right->GetLeftHight() > pCurNode->right->GetRightHight())
		{
			RL_Rotate(pCurNode, pParent, nLRFlag);  // RL型 平衡处理
		}
		else
			RR_Rotate(pCurNode, pParent, nLRFlag);  // RR型 平衡处理
	}
}

BOOL CAVLTree::LL_Rotate( TreeNode *& pCurNode, TreeNode *& pParentNode, const int &nLRFlag )
{
	cout << "Call LL_Rotate() succeeded!\n";
	TreeNode *pTempRoot = pCurNode->left;
	TreeNode *pTempRootR = pTempRoot->right;
	pTempRoot->right = pCurNode;
	pCurNode->left = pTempRootR;
	if (pParentNode)
	{
		if (nLRFlag == 0)
		{
			pParentNode->left = pTempRoot;
		}
		else
			pParentNode->right = pTempRoot;
	}
	else  // root
		m_root = pTempRoot;
	
	return TRUE;
}

BOOL CAVLTree::LR_Rotate( TreeNode *& pCurNode, TreeNode *& pParentNode, const int &nLRFlag )
{
	cout << "Call LR_Rotate() run succeeded!\n";
	TreeNode *pNodeA = pCurNode->left;
	TreeNode *pNodeB = pNodeA->right;
	TreeNode *pPreB = pNodeA;
	TreeNode *pTempRoot = NULL;
	while(pNodeB->right)
	{
		pPreB = pNodeB;
		pNodeB = pNodeB->right;
	}
	pTempRoot = pNodeB;
	if (pTempRoot == NULL)   // pNodeA 的右子树为空
	{
		pTempRoot = pNodeA;
		pTempRoot->right = pCurNode;
		pCurNode->left = NULL;
	}
	else
	{
		pPreB->right = pTempRoot->left;
		pTempRoot->left = pNodeA;
		pTempRoot->right = pCurNode;
		pCurNode->left = NULL;
	}
	if (pParentNode)
	{
		if (nLRFlag == 0)
		{
			pParentNode->left = pTempRoot;
		}
		else
			pParentNode->right = pTempRoot;
	}
	else  // root
		m_root = pTempRoot;

	
	return TRUE;
}

BOOL CAVLTree::RR_Rotate( TreeNode *& pCurNode, TreeNode *& pParentNode, const int &nLRFlag )
{
	cout << "Call RR_Rotate() succeeded!\n";
	TreeNode *pTempRoot = pCurNode->right;
	TreeNode *pTempRootL = pTempRoot->left;
	pTempRoot->left = pCurNode;
	pCurNode->right = pTempRootL;
	if (pParentNode)
	{
		if (nLRFlag == 0)
		{
			pParentNode->left = pTempRoot;
		}
		else
			pParentNode->right = pTempRoot;
	}
	else  // root
		m_root = pTempRoot;
	
	return TRUE;
}

BOOL CAVLTree::RL_Rotate( TreeNode *& pCurNode, TreeNode *& pParentNode, const int &nLRFlag )
{
	cout << "Call RL_Rotate() run succeeded!\n";
	TreeNode *pNodeA = pCurNode->right;
	TreeNode *pNodeB = pNodeA->left;
	TreeNode *pPreB = pNodeA;
	while(pNodeB->left)
	{
		pPreB = pNodeB;
		pNodeB = pPreB->left;
	}
	TreeNode *pTempRoot = pNodeB;
	if (pTempRoot == NULL)
	{
		pTempRoot = pNodeA;
		pCurNode->right = NULL;
		pTempRoot->left = pCurNode;
	}
	else
	{
		pTempRoot->left = pCurNode;
		pCurNode->right = NULL;
		pPreB->left = pTempRoot->right;
		pTempRoot->right = pNodeA;
	}

	if (pParentNode)
	{
		if (nLRFlag == 0)
		{
			pParentNode->left = pTempRoot;
		}
		else
			pParentNode->right = pTempRoot;
	}
	else  // root
		m_root = pTempRoot;
	
	
	return TRUE;
}

void CAVLTree::PrintTree()   // 打印树结点
{
	if (!m_root)
	{
		cout << "The tree is empty!\n";
		return ;
	}
	PrintTree(m_root);  // 调用递归遍历结点函数
	cout << endl;
}

void CAVLTree::PrintTree( TreeNode *& pNode )
{
	if (pNode)
	{
		cout << pNode->data << "  ";
		PrintTree(pNode->left);
		PrintTree(pNode->right);
	}
}

void CAVLTree::EmptyTree()  // 清空结点，释放结点空间
{
	if (!m_root)
	{
		return;
	}
	CStack stNodes(10000);
	GetNodes(stNodes, m_root);
	TreeNode *pTemp;
	if (stNodes.Pop(pTemp))
	{
		delete pTemp;
	}
	m_root = NULL;
	cout << "Emptied the tree succeeded!\n";
}

void CAVLTree::GetNodes( CStack &st, TreeNode *pNode )
{
	if (pNode)
	{
		GetNodes(st, pNode->left);
		GetNodes(st, pNode->right);
	}
}

void CAVLTree::SwapNodes( TreeNode *&pNode1, TreeNode *&pNode2 )  // 交换两个结点
{
	ElemType val = pNode1->data;
	pNode1->data = pNode2->data;
	pNode2->data = val;
}

测试文件

/// test.cpp for test the class CAVLTree
#include <iostream>
#include "AVLTree.h"
#include <ctime>
using namespace std;

void LL_Test();
void RR_Test();
void LR_Test();
void RL_Test();
void time_Test();
void main()
{
	time_Test();
// 	LL_Test();
// 	RR_Test();
//	LR_Test();
//	RL_Test();
	return;
}

void time_Test()
{
	cout << "////////////////time_Test Function.\n";
	CAVLTree tree;
	clock_t t_insert, t_print, t_find;  // 分别记录插入数据时间和打印数据时间
	int nArrNodes[30000];
	srand((unsigned)time(NULL));
//	srand(0);
	t_insert = clock();
	for (int i = 0; i < 30000; i++)
	{
		nArrNodes[i] = rand();
		while (1)
		{
			if (tree.Find(nArrNodes[i]) == NULL)
			{
				tree.Insert(nArrNodes[i]);
				break;
			}
			else
				nArrNodes[i] = rand();
		}	
	}
	t_insert = clock() - t_insert;
	t_print = clock();
	tree.PrintTree();
	t_print = clock() - t_print;
	t_find = clock();
	tree.Find(32725);
	t_find = clock() - t_find;
	cout << "The tree nodes count is: " << tree.GetSize() << endl;
	cout << "Insert data need time: " << t_insert << "ms." << endl;
	cout << "Print data need time: " << t_print << "ms." << endl;
	cout << "Find one data need time: " << t_find << "ms." << endl;
}

void LL_Test()
{
	cout << "/////////////////LL_Test" << endl;
	CAVLTree tree;
	tree.Insert(9);
	tree.Insert(5);
	tree.Insert(11);
	tree.Insert(6);
	tree.Insert(15);
	tree.Insert(4);
	tree.Insert(3);
	tree.PrintTree();
	tree.Remove(9);
	tree.PrintTree();
	cout << endl;
	return;
}

void RR_Test()
{
	cout << "//////////////////RR_Test" << endl;
	CAVLTree tree;
	tree.Insert(4);
	tree.Insert(2);
	tree.Insert(6);
	tree.Insert(5);
	tree.Insert(9);
	tree.Insert(11);
	tree.PrintTree();
	cout << endl;
	return;
}

void LR_Test()
{
	cout << "////////////////LR_Test" << endl;
	CAVLTree tree;
	tree.Insert(9);
	tree.Insert(4);
	tree.Insert(11);
	tree.Insert(2);
	tree.Insert(6);
	tree.Insert(5);
	tree.PrintTree();
	tree.Remove(6);
	tree.PrintTree();
	cout << endl;
	return;
}

void RL_Test()
{
	cout << "//////////////////RL_Test" << endl;
	CAVLTree tree;
	tree.Insert(4);
	tree.Insert(2);
	tree.Insert(9);
	tree.Insert(6);
	tree.Insert(11);
	tree.Insert(5);
	tree.PrintTree();
	tree.Insert(7);
	tree.Insert(8);
	tree.PrintTree();
	tree.Remove(6);
	tree.PrintTree();
	tree.Remove(11);
	tree.PrintTree();
	tree.Remove(7);
	tree.PrintTree();
	tree.Remove(8);
	tree.PrintTree();
	tree.Remove(9);
	tree.PrintTree();
	tree.Remove(4);
	tree.PrintTree();
	tree.Remove(2);
	tree.PrintTree();
	tree.Remove(5);
	tree.PrintTree();
	cout << endl;
	return;
}