数据结构实现 10.1：AVL树（C++版）

1. 概念及基本框架

AVL树 （即自平衡树）是一种二分搜索树，AVL 树要求每个结点的左右子树的树高度差不超过 1 。

因为在二分搜索树中，不管二叉树左右子树的高度，所以二分搜索树的最坏情况可以退化成链表。所以这里我们需要对二分搜索树进行适当的修正，提高效率。
与二分搜索树类似，我们先给出AVL树的结点类：

template <typename K, typename V>
class AVLNode{
public:
	AVLNode(K key, V value){
		this->key = key;
		this->value = value;
		height = 1;
		left = NULL;
		right = NULL;
	}
public:
	K key;
	V value;
	int height;
	AVLNode<K, V> *left;
	AVLNode<K, V> *right;
};

这里我们增加了高度变量，用来记录每个结点的当前高度。
下面给出AVL树大体框架：

template <typename K, typename V>
class AVLTree{
public:
	AVLTree(){
		root = NULL;
		m_size = 0;
	}
	...
private:
	AVLNode<K, V> *root;
	int m_size;
};

与二分搜索树十分类似，所以这里不再赘述。
接下来我们就对AVL树的增加、删除、查找、遍历以及一些其他基本操作用代码去实现。

2. 基本操作程序实现

2.1 增加操作

既然AVL树是二分搜索树的一种，所以其基本实现代码与二分搜索树类似，我们直接给出：

template <typename K, typename V>
class AVLTree{
public:
	...
	//增加操作
	void add(K key, V value);
	...
private:
	AVLNode<K, V>* add(AVLNode<K, V>* node, K key, V value);
	...
};

类外实现代码：

template <typename K, typename V>
void AVLTree<K, V>::add(K key, V value){
	root = add(root, key, value);
}
template <typename K, typename V>
AVLNode<K, V>* AVLTree<K, V>::add(AVLNode<K, V>* node, K key, V value){
	if (node == NULL){
		m_size++;
		return new AVLNode<K, V>(key, value);
	}
	else if (key < node->key){
		node->left = add(node->left, key, value);
	}
	else if (key > node->key){
		node->right = add(node->right, key, value);
	}
	else if (key == node->key){
		node->value = value;
		return node;
	}
	//计算平衡因子
	node->height = max(getHeight(node->left), getHeight(node->right)) + 1;
	int balanceFactor = getBalanceFactor(node);
	//平衡维护
	if (balanceFactor > 1 && getBalanceFactor(node->left) >= 0){
		return rightRotate(node);
	}
	else if (balanceFactor < -1 && getBalanceFactor(node->right) <= 0){
		return leftRotate(node);
	}
	else if (balanceFactor > 1 && getBalanceFactor(node->left) < 0){
		node->left = leftRotate(node->left);
		return rightRotate(node);
	}
	else if (balanceFactor < -1 && getBalanceFactor(node->right) > 0){
		node->right = rightRotate(node->right);
		return leftRotate(node);
	}
	return node;
}

这里需要注意的是在进行完增加操作之后我们需要进行一定的修正。修正之前需要得到一些基础信息，用下面几个函数来实现。
max 返回其中的最大值。
getHeight 返回一个结点的高度。
getBalanceFactor 返回一个结点的平衡因子。（用左子树高度减去右子树高度）
leftRotate、rightRotate 分别表示左旋、右旋，具体实现方法后面详细讲述。
前面三个函数代码如下：

template <typename K, typename V>
class AVLTree{
	...
private:
	...
	int max(int a, int b){ return a > b ? a : b; }
	...
};

template <typename K, typename V>
int AVLTree<K, V>::getHeight(AVLNode<K, V>* node){
	if (node == NULL){
		return 0;
	}
	return node->height;
}
template <typename K, typename V>
int AVLTree<K, V>::getBalanceFactor(AVLNode<K, V>* node){
	int lHeight = 0, rHeight = 0;
	if (node == NULL){
		return 0;
	}
	if (node->left != NULL){
		lHeight = node->left->height;
	}
	if (node->right != NULL){
		rHeight = node->right->height;
	}
	return lHeight - rHeight;
}

接下来我们详细讲述修正的左旋和右旋函数。
首先我们分情况，当一个新增加的结点需要修正，一定是其祖父结点一脉单传到了新结点。那么，作为一个新增加的结点，有两种可能，在父结点的左/右，而父结点也可以在祖父节点的左/右，所以一共有四种。我们可以简称为左左、左右、右左、右右四种情况。

2.1.1 左左

注：这里每个结点都挂有子树（子树可以为空），这些子树本身已经是平衡的了（下同，这里考虑的是修正自下而上递推到这里），我们可以忽略。
如上图，结点 C 左右子树平衡因子为 2 > 1 ，所以需要修正，修正过程如下：
1.新建一个临时结点 P 指向 B 的右子结点。
2.把 C 作为 B 的右子结点。
3.把 P 作为 C 的左子结点。
4.将结点 B 返回。
经过旋转，结点的数据依旧满足二分搜索树的基本性质，所以可以旋转，这里结点 B 和结点 C 发生了 右旋 。
图中虚线圈住的结点表示实际发生了旋转的结点，橙色（紫色）的线表示发生变化的支，黄色的线表示返回结点的指向，后面不再赘述。
右旋代码实现如下：

template <typename K, typename V>
AVLNode<K, V>* AVLTree<K, V>::rightRotate(AVLNode<K, V>* node){
	AVLNode<K, V>* res = node->left;
	AVLNode<K, V>* p = res->right;
	res->right = node;
	node->left = p;
	//维护平衡
	node->height = max(getHeight(node->left), getHeight(node->right)) + 1;
	res->height = max(getHeight(res->left), getHeight(res->right)) + 1;
	return res;
}

2.1.2 右右

如上图，结点 A 左右子树平衡因子为 | - 2 | > 1 ，所以需要修正，修正过程如下：
1.新建一个临时结点 P 指向 B 的左子结点。
2.把 A 作为 B 的左子结点。
3.把 P 作为 A 的右子结点。
4.将结点 B 返回。
经过旋转，结点的数据依旧满足二分搜索树的基本性质，所以可以旋转，这里结点 A 和结点 B 发生了 左旋 。
下面分析两种比较复杂的情况，左右和右左。
左旋代码实现如下：

template <typename K, typename V>
AVLNode<K, V>* AVLTree<K, V>::leftRotate(AVLNode<K, V>* node){
	AVLNode<K, V>* res = node->right;
	AVLNode<K, V>* p = res->left;
	res->left = node;
	node->right = p;
	//维护平衡
	node->height = max(getHeight(node->left), getHeight(node->right)) + 1;
	res->height = max(getHeight(res->left), getHeight(res->right)) + 1;
	return res;
}

2.1.3 左右

如上图，结点 C 左右子树平衡因子为 2 > 1 ，所以需要修正，修正过程如下：
1.把结点 A 和结点 B 左旋。（这时可看作右右情况）
2.把结点 B 和结点 C 右旋。（这时是左左情况）
3.返回结点 B 。
经过旋转，结点的数据依旧满足二分搜索树的基本性质。

2.1.4 右左

如上图，结点 A 左右子树平衡因子为 | - 2 | > 1 ，所以需要修正，修正过程如下：
1.把结点 B 和结点 C 右旋。（这时可看作左左情况）
2.把结点 A 和结点 B 左旋。（这时是右右情况）
3.返回结点 B 。
经过旋转，结点的数据依旧满足二分搜索树的基本性质。

2.1.5 增加操作修正总结

1.左左：右旋祖父结点和父结点，返回父结点。
2.右右：左旋祖父结点和父结点，返回父结点。
3.左右：左旋父结点和子结点，右旋祖父结点和新的父结点（即原来的子结点），返回新的父结点（即原来的子结点）。
4.右左：右旋父结点和子结点，左旋祖父结点和新的父结点（即原来的子结点），返回新的父结点（即原来的子结点）。
左左右，右右左，兄做侄，祖成兄，父替祖，返林中。
左右左，成左左，右左右，成右右，人虽新，法依旧。

2.2 删除操作

删除操作与二分搜索树类似，只是增加修正过程：

template <typename K, typename V>
class AVLTree{
public:
	...
	//删除操作
	void remove(K key);
	...
private:
	...
	AVLNode<K, V>* remove(AVLNode<K, V>* node, K key);
	...
};

类外实现如下：

template <typename K, typename V>
void AVLTree<K, V>::remove(K key){
	root = remove(root, key);
}
template <typename K, typename V>
AVLNode<K, V>* AVLTree<K, V>::remove(AVLNode<K, V> *node, K key){
	if (node == NULL){
		return node;
	}
	AVLNode<K, V> *resNode = node;
	if (key < node->key){
		node->left = remove(node->left, key);
	}
	else if (key > node->key){
		node->right = remove(node->right, key);
	}
	else if (key == node->key){
		if (node->left == NULL){
			AVLNode<K, V> *rightNode = node->right;
			delete node;
			m_size--;
			resNode = rightNode;
		}
		else if (node->right == NULL){
			AVLNode<K, V> *leftNode = node->left;
			delete node;
			m_size--;
			resNode = leftNode;
		}
		else{
			AVLNode<K, V> *minNode = node->right;
			for (; minNode->left; minNode = minNode->left);
			minNode->right = remove(node->right, minNode->key);
			minNode->left = node->left;
			node->left = node->right = NULL;
			resNode = minNode;
		}
	}
	if (resNode == NULL){
		return NULL;
	}
	//计算平衡因子
	resNode->height = max(getHeight(resNode->left), getHeight(resNode->right)) + 1;
	int balanceFactor = getBalanceFactor(resNode);
	//平衡维护
	if (balanceFactor > 1 && getBalanceFactor(resNode->left) >= 0){
		return rightRotate(resNode);
	}
	else if (balanceFactor < -1 && getBalanceFactor(resNode->right) <= 0){
		return leftRotate(resNode);
	}
	else if (balanceFactor  > 1 && getBalanceFactor(resNode->left) < 0){
		resNode->left = leftRotate(resNode->left);
		return rightRotate(resNode);
	}
	else if (balanceFactor < -1 && getBalanceFactor(resNode->right) > 0){
		resNode->right = rightRotate(resNode->right);
		return leftRotate(resNode);
	}
	return resNode;
}

2.2.1 左右均空

2.2.2 左空右不空

2.2.3 左不空右空

2.2.4 左右均不空

2.2.5 删除操作修正总结

1.左右均空：直接删除该结点，返回空。
2.左空右不空：删除该结点，返回右子结点。
3.左不空右空：删除该结点，返回左子结点。
4.左右均不空：找到该结点的右子结点的最左结点，即该结点的直接后继，然后将该结点的值赋给要删除的结点，然后删掉这个后继结点。
左子空，右子从，右子空，左子从，若无子，返作空。
寻右子，左子孙，无左子，方可停，假作真，除真身。
删除后需要在删除结点处开始维护平衡，方法同增加操作时的修正相同，这里不再赘述。

2.3 查找操作

这里我们提供两个查询函数，contains 和 get ，函数实现代码如下。

template <typename K, typename V>
class AVLTree{
public:
	...
	//查找操作
	bool contains(K key);
	V get(K key);
	...
private:
	...
	bool contains(AVLNode<K, V>* node, K key);
	...
};

类外实现如下：

template <typename K, typename V>
bool AVLTree<K, V>::contains(K key){
	return contains(root, key);
}
template <typename K, typename V>
bool AVLTree<K, V>::contains(AVLNode<K, V> *node, K key){
	if (node == NULL){
		return false;
	}
	if (key == node->key){
		return true;
	}
	else if (key < node->key){
		return contains(node->left, key);
	}
	else{
		return contains(node->right, key);
	}
}
template <typename K, typename V>
V AVLTree<K, V>::get(K key){
	AVLNode<K, V> *node = root;
	while (node != NULL){
		if (key == node->key){
			return node->value;
		}
		else if (key < node->key){
			node = node->left;
		}
		else if (key > node->key){
			node = node->right;
		}
	}
	cout << "查询失败！AVL树中不存在" << key << "！" << endl;
	return NULL;
}

2.4 遍历操作

这里与二分搜索树类似，所以只给出中序遍历代码，其余不再赘述。

template <typename K, typename V>
class AVLTree{
public:
	...
	//遍历操作
	void inOrder();
private:
	...
	void inOrder(AVLNode<K, V> *node);
	...
};

类外实现如下：

template <typename K, typename V>
void AVLTree<K, V>::inOrder(){
	cout << "AVL Tree: " << "Size = " << m_size << endl;
	cout << "中序遍历：";
	inOrder(root);
	cout << endl;
}
template <typename K, typename V>
void AVLTree<K, V>::inOrder(AVLNode<K, V> *node){
	if (node == NULL){
		return;
	}
	inOrder(node->left);
	cout << node->key << "  ";
	inOrder(node->right);
}

2.5 其他操作

AVL树还有一些其他的操作，包括 AVL树大小 等的查询操作。

template <typename K, typename V>
class AVLTree{
public:
	...
	int size(){
		return m_size;
	}
	bool isEmpty(){
		return m_size == 0;
	}
	...
};

3. 算法复杂度分析

3.1 增加操作

函数	最坏复杂度	平均复杂度
add	O(logn)	O(logn)

3.2 删除操作

函数	最坏复杂度	平均复杂度
remove	O(logn)	O(logn)

3.3 查找操作

函数	最坏复杂度	平均复杂度
contains	O(logn)	O(logn)
get	O(logn)	O(logn)

总体情况：

操作	时间复杂度
增	O(logn)
删	O(logn)
查	O(logn)

4. 完整代码

程序完整代码（这里使用了头文件的形式来实现类）如下：

#ifndef __AVLTREE_H__
#define __AVLTREE_H__

template <typename K, typename V>
class AVLNode{
public:
	AVLNode(K key, V value){
		this->key = key;
		this->value = value;
		height = 1;
		left = NULL;
		right = NULL;
	}
public:
	K key;
	V value;
	int height;
	AVLNode<K, V> *left;
	AVLNode<K, V> *right;
};

template <typename K, typename V>
class AVLTree{
public:
	AVLTree(){
		root = NULL;
		m_size = 0;
	}
	int size(){
		return m_size;
	}
	bool isEmpty(){
		return m_size == 0;
	}
	//增加操作
	void add(K key, V value);
	//删除操作
	void remove(K key);
	//查找操作
	bool contains(K key);
	V get(K key);
	//遍历操作
	void inOrder();
private:
	AVLNode<K, V>* add(AVLNode<K, V>* node, K key, V value);
	AVLNode<K, V>* remove(AVLNode<K, V>* node, K key);
	bool contains(AVLNode<K, V>* node, K key);
	int getHeight(AVLNode<K, V>* node);
	int max(int a, int b){ return a > b ? a : b; }
	int getBalanceFactor(AVLNode<K, V>* node);
	AVLNode<K, V>* rightRotate(AVLNode<K, V>* node);
	AVLNode<K, V>* leftRotate(AVLNode<K, V>* node);
	void inOrder(AVLNode<K, V> *node);
private:
	AVLNode<K, V> *root;
	int m_size;
};

template <typename K, typename V>
void AVLTree<K, V>::add(K key, V value){
	root = add(root, key, value);
}
template <typename K, typename V>
AVLNode<K, V>* AVLTree<K, V>::add(AVLNode<K, V>* node, K key, V value){
	if (node == NULL){
		m_size++;
		return new AVLNode<K, V>(key, value);
	}
	else if (key < node->key){
		node->left = add(node->left, key, value);
	}
	else if (key > node->key){
		node->right = add(node->right, key, value);
	}
	else if (key == node->key){
		node->value = value;
		return node;
	}
	//计算平衡因子
	node->height = max(getHeight(node->left), getHeight(node->right)) + 1;
	int balanceFactor = getBalanceFactor(node);
	//平衡维护
	if (balanceFactor > 1 && getBalanceFactor(node->left) >= 0){
		return rightRotate(node);
	}
	else if (balanceFactor < -1 && getBalanceFactor(node->right) <= 0){
		return leftRotate(node);
	}
	else if (balanceFactor > 1 && getBalanceFactor(node->left) < 0){
		node->left = leftRotate(node->left);
		return rightRotate(node);
	}
	else if (balanceFactor < -1 && getBalanceFactor(node->right) > 0){
		node->right = rightRotate(node->right);
		return leftRotate(node);
	}
	return node;
}

template <typename K, typename V>
void AVLTree<K, V>::remove(K key){
	root = remove(root, key);
}
template <typename K, typename V>
AVLNode<K, V>* AVLTree<K, V>::remove(AVLNode<K, V> *node, K key){
	if (node == NULL){
		return node;
	}
	AVLNode<K, V> *resNode = node;
	if (key < node->key){
		node->left = remove(node->left, key);
	}
	else if (key > node->key){
		node->right = remove(node->right, key);
	}
	else if (key == node->key){
		if (node->left == NULL){
			AVLNode<K, V> *rightNode = node->right;
			delete node;
			m_size--;
			resNode = rightNode;
		}
		else if (node->right == NULL){
			AVLNode<K, V> *leftNode = node->left;
			delete node;
			m_size--;
			resNode = leftNode;
		}
		else{
			AVLNode<K, V> *minNode = node->right;
			for (; minNode->left; minNode = minNode->left);
			minNode->right = remove(node->right, minNode->key);
			minNode->left = node->left;
			node->left = node->right = NULL;
			resNode = minNode;
		}
	}
	if (resNode == NULL){
		return NULL;
	}
	//计算平衡因子
	resNode->height = max(getHeight(resNode->left), getHeight(resNode->right)) + 1;
	int balanceFactor = getBalanceFactor(resNode);
	//平衡维护
	if (balanceFactor > 1 && getBalanceFactor(resNode->left) >= 0){
		return rightRotate(resNode);
	}
	else if (balanceFactor < -1 && getBalanceFactor(resNode->right) <= 0){
		return leftRotate(resNode);
	}
	else if (balanceFactor  > 1 && getBalanceFactor(resNode->left) < 0){
		resNode->left = leftRotate(resNode->left);
		return rightRotate(resNode);
	}
	else if (balanceFactor < -1 && getBalanceFactor(resNode->right) > 0){
		resNode->right = rightRotate(resNode->right);
		return leftRotate(resNode);
	}
	return resNode;
}

template <typename K, typename V>
bool AVLTree<K, V>::contains(K key){
	return contains(root, key);
}
template <typename K, typename V>
bool AVLTree<K, V>::contains(AVLNode<K, V> *node, K key){
	if (node == NULL){
		return false;
	}
	if (key == node->key){
		return true;
	}
	else if (key < node->key){
		return contains(node->left, key);
	}
	else{
		return contains(node->right, key);
	}
}
template <typename K, typename V>
V AVLTree<K, V>::get(K key){
	AVLNode<K, V> *node = root;
	while (node != NULL){
		if (key == node->key){
			return node->value;
		}
		else if (key < node->key){
			node = node->left;
		}
		else if (key > node->key){
			node = node->right;
		}
	}
	cout << "查询失败！AVL树中不存在" << key << "！" << endl;
	return NULL;
}

template <typename K, typename V>
int AVLTree<K, V>::getHeight(AVLNode<K, V>* node){
	if (node == NULL){
		return 0;
	}
	return node->height;
}
template <typename K, typename V>
int AVLTree<K, V>::getBalanceFactor(AVLNode<K, V>* node){
	int lHeight = 0, rHeight = 0;
	if (node == NULL){
		return 0;
	}
	if (node->left != NULL){
		lHeight = node->left->height;
	}
	if (node->right != NULL){
		rHeight = node->right->height;
	}
	return lHeight - rHeight;
}
template <typename K, typename V>
AVLNode<K, V>* AVLTree<K, V>::rightRotate(AVLNode<K, V>* node){
	AVLNode<K, V>* res = node->left;
	AVLNode<K, V>* p = res->right;
	res->right = node;
	node->left = p;
	//维护平衡
	node->height = max(getHeight(node->left), getHeight(node->right)) + 1;
	res->height = max(getHeight(res->left), getHeight(res->right)) + 1;
	return res;
}
template <typename K, typename V>
AVLNode<K, V>* AVLTree<K, V>::leftRotate(AVLNode<K, V>* node){
	AVLNode<K, V>* res = node->right;
	AVLNode<K, V>* p = res->left;
	res->left = node;
	node->right = p;
	//维护平衡
	node->height = max(getHeight(node->left), getHeight(node->right)) + 1;
	res->height = max(getHeight(res->left), getHeight(res->right)) + 1;
	return res;
}

template <typename K, typename V>
void AVLTree<K, V>::inOrder(){
	cout << "AVL Tree: " << "Size = " << m_size << endl;
	cout << "中序遍历：";
	inOrder(root);
	cout << endl;
}
template <typename K, typename V>
void AVLTree<K, V>::inOrder(AVLNode<K, V> *node){
	if (node == NULL){
		return;
	}
	inOrder(node->left);
	cout << node->key << "  ";
	inOrder(node->right);
}

#endif