AVL树的实现

AVL树是带有平衡条件的二叉查找树，每个节点的左子树和右子树的高度最多差1。

#include <iostream>

template <typename T>
class AVLTree
{
public:
	AVLTree();
	AVLTree(const AVLTree & at);
	~AVLTree();
	
	//插入
	void insert(const T & x);
	//删除
	void remove(const T & x);
	//是否包含（查找）
	bool contains(const T & x) const;
	//最大值
	const T & findMax() const;
	//最小值
	const T & findMin() const;
	//判空
	bool isEmpty() const;
	//清空树
	void makeEmpty();
	//打印树
	void print();

	//左右高度差超过 1 就视为不平衡
	static const int ALLOWED_IMBALANCE = 1;

private:
	struct AvlNode
	{
		T element;		//节点值
		AvlNode *left;	//左节点
		AvlNode *right;	//右节点
		int height;		//节点高度

		AvlNode(const T & elem, AvlNode *l, AvlNode *r, int h=0) 
			: element{elem}, left{l}, right{r}, height{h} { }
	};
	
	AvlNode *root;	//根节点

	int height(AvlNode * & t) const;	//获取节点高度

	AvlNode * clone(AvlNode * & t) const;
	void insert(const T & x, AvlNode * & t);
	void remove(const T & x, AvlNode * & t);
	bool contains(const T & x, AvlNode * & t) const;
	AvlNode * findMax(AvlNode * & t) const;
	AvlNode * findMin(AvlNode * & t) const;
	void makeEmpty(AvlNode * & t);
	void print(AvlNode * & t);

	//四种不平衡的情况
	void roteWithLeftChild(AvlNode * & k2);
	void roteWithRightChild(AvlNode * & k2);
	void doubleRoteWithLeftChild(AvlNode * & k3);
	void doubleRoteWithRightChild(AvlNode * & k3);

	//设置平衡
	void balance(AvlNode * & t);
};

/**
* 成员函数实现
*/

//无参构造函数
template <typename T>
AVLTree<T>::AVLTree() : root{ nullptr }
{

}

//拷贝构造函数
template <typename T>
AVLTree<T>::AVLTree(const AVLTree & at) : root{ nullptr }
{
	root = clone(at.root);
}

template <typename T>
typename AVLTree<T>::AvlNode * AVLTree<T>::clone(AvlNode * & t) const
{
	if (t == nullptr)
		return nullptr;
	else
		return new AvlNode(t->element, clone(t->left), clone(t->right));
}

//析构函数
template <typename T>
AVLTree<T>::~AVLTree()
{
	makeEmpty();
}

template <typename T>
void AVLTree<T>::makeEmpty()
{
	makeEmpty(root);
}

template <typename T>
void AVLTree<T>::makeEmpty(AvlNode * & t)
{
	if (t != nullptr)
	{
		makeEmpty(t->left);
		makeEmpty(t->right);
		delete t;
	}
	t = nullptr;
}

/**
* 返回节点t的高度,nullptr则返回-1
*/
template <typename T>
int AVLTree<T>::height(AvlNode * & t) const
{
	if (t == nullptr)
		return 0;
	int L = height(t->left);
	int R = height(t->right);
	return L < R ? L+1 : R+1;
}

//获取较大值
auto max = [](int x, int y)
{
	return x > y ? x : y;
};

/**
* 外侧插入，左侧高（单旋转）
*			k2                     k1
*		   /  \                   / \
*		  k1   z      -->       x     k2
*		/  \                   /      / \
*	   x    y                        y   z
*     /
*/
template <typename T>
void AVLTree<T>::roteWithLeftChild(AvlNode * & k2)
{
	AvlNode *k1 = k2->left;
	k2->left = k1->right;
	k1->right = k2;

	//旋转后重置高度
	k2->height = max(height(k2->left), height(k2->right)) + 1;
	k1->height = max(height(k1->left), k2->height) + 1;

	k2 = k1; //设置新根
}

/**
* 外侧插入，右侧高（单旋转）
*			k1                     k2
*		   /  \                   / \
*		  k2   z      <--       x     k1
*		/  \    \                     / \
*	   x    y                        y   z
*                                         \
*/
template <typename T>
void AVLTree<T>::roteWithRightChild(AvlNode * & k2)
{
	AvlNode *k1 = k2->right;
	k2->right = k1->left;
	k1->left = k2;

	//旋转后重置高度
	k2->height = max(height(k2->left), height(k2->right)) + 1;
	k1->height = max(height(k1->left), k2->height) + 1;

	k2 = k1; //设置新根
}

/**
* 内侧插入，左高（双旋转）
*		   k3                  k3                  k2
*	      / \                 /  \                /  \
*        k1  z      -->      k2   z      -->     k1    k3
*  		/ \                 /                   / \      \
*	   x   k2               k1                  x   y      z
*	      /                / \
*		 y                x   y
*/
template <typename T>
void AVLTree<T>::doubleRoteWithLeftChild(AvlNode * & k3)
{
	roteWithRightChild(k3->left);
	roteWithLeftChild(k3);
}

/**
* 内侧插入，右高（双旋转）
*        k3                         k3                     k2      
*       /  \                       /  \                   /  \
*      z    k1                    z    k2                k3    k1
*           / \         -->              \       -->    /      / \
*          k2   x                         k1           z      y   x
*           \                             / \
*            y                           y   x 
*/
template <typename T>
void AVLTree<T>::doubleRoteWithRightChild(AvlNode * & k3)
{
	roteWithLeftChild(k3->right);
	roteWithRightChild(k3);
}

//设置平衡
template <typename T>
void AVLTree<T>::balance(AvlNode * & t)
{
	if (t == nullptr)
		return;

	if (height(t->left) - height(t->right) > ALLOWED_IMBALANCE)	//左侧高
	{
		if (height(t->left->left) >= height(t->left->right))	//外侧插入
			roteWithLeftChild(t);
		else  //内侧插入
			doubleRoteWithLeftChild(t);
	}
	else if (height(t->right) - height(t->left) > ALLOWED_IMBALANCE)  //右侧高
	{
		if (height(t->right->right) >= height(t->right->left)) //外侧插入
			roteWithRightChild(t);
		else  //内侧插入
			doubleRoteWithRightChild(t);
	}

	//重置高
	t->height = max(height(t->left), height(t->right)) + 1;
}

//插入数据
template <typename T>
void AVLTree<T>::insert(const T & x)
{
	insert(x, root);
}

template <typename T>
void AVLTree<T>::insert(const T & x, AvlNode * & t)
{
	if (t == nullptr)
		t = new AvlNode(x, nullptr, nullptr);
	else if (x < t->element)
		insert(x, t->left);
	else if (x > t->element)
		insert(x, t->right);

	balance(t);
}

//删除数据
template <typename T>
void AVLTree<T>::remove(const T & x)
{
	remove(x, root);
}

template <typename T>
void AVLTree<T>::remove(const T & x, AvlNode * & t)
{
	if (t == nullptr)
		return;
	
	if (x < t->element)
		remove(x, t->left);
	else if (x > t->element)
		remove(t->left);
	else if ((t->left != nullptr) && (t->right != nullptr))
	{
		t->element = findMin(t->right)->element;
		remove(t->element, t->right);
	}
	else
	{
		AvlNode *oldNode = t;
		t = (t->left != nullptr) ? t->left : t->right;
		delete oldNode;
	}

	balance(t);
}

//查找
template <typename T>
bool AVLTree<T>::contains(const T & x) const
{
	return contains(x, root);
}

template <typename T>
bool AVLTree<T>::contains(const T & x, AvlNode * & t) const
{
	if (t == nullptr)
		return false;

	if (x < t->element)
		contains(x, t->left);
	else if (x > t->element)
		contains(x, t->right);
	else
		return true;
}

//判空
template <typename T>
bool AVLTree<T>::isEmpty() const
{
	if (root == nullptr)
		return true;
	return false;
}

//查找最大值
template <typename T>
const T & AVLTree<T>::findMax() const
{
	return findMax(root)->element;
}

template <typename T>
typename AVLTree<T>::AvlNode * AVLTree<T>::findMax(AvlNode * & t) const
{
	if (t != nullptr)
		while (t->right != nullptr)
			t = t->right;
	return t;
}

//查找最小值
template <typename T>
const T & AVLTree<T>::findMin() const
{
	return findMin(root)->element;
}

template <typename T>
typename AVLTree<T>::AvlNode * AVLTree<T>::findMin(AvlNode * & t) const
{
	if (t != nullptr)
		while (t->left != nullptr)
			t = t->left;
	return t;
}

//打印树
template <typename T>
void AVLTree<T>::print()
{
	if (isEmpty())
		std::cout << "空树" << std::endl;
	else
		print(root);
}

template <typename T>
void AVLTree<T>::print(AvlNode * & t)
{
	if (t != nullptr)
	{
		print(t->left);
		std::cout << t->element << "  ";
		print(t->right);
	}
}