平衡二叉树实现(Balanced Binary Tree)

平衡二叉树实现(Balanced Binary Tree)

二叉搜索树查东西很快,但它有个致命的脾气:如果你按从小到大的顺序把数据插进去,树就会退化成一条链表。查一个值从原来的每层砍一半变成从头到尾挨个比,复杂度直接从 O(log n) 掉到 O(n)。工程上不能接受这种看运气的性能。平衡二叉树就是为了治这个病来的。

平衡二叉树也叫 AVL 树,1962 年由 Adelson-Velsky 和 Landis 提出。它给每个节点加了一条硬规矩,任何节点的左右子树高度差不能超过 1。一旦插入或删除打破这条规矩,就通过旋转把树掰回平衡。这一节我们讲清楚为什么要平衡、高度和平衡因子怎么算、四种旋转怎么转,再配上完整的 C++ 实现。读之前得先会写普通的二叉搜索树,这一节的大部分操作都和它一样,差别只在插入删除之后多了一步调平衡。

1. 平衡解决什么

我们先看清楚不看病在哪。二叉搜索树的查找效率取决于树高。理想情况下树是匀称的,n 个节点高约 log2 n,查找每层走一条路,复杂度 O(log n)。但如果你依次插入 1、2、3、4、5,每个新节点都比前一个大,全挂在右孩子上,树就拉成一条右斜链,高度变成 n,查找退化成 O(n)。

AVL 树的思路是,不管你按什么顺序插,插完都检查一下,哪边太重了就转一下,让树始终保持匀称。这样最好情况和最坏情况下树高都压在 O(log n),插入、查找、删除稳定在 O(log n)。代价是插入和删除时偶尔要旋转,多花一点时间,但换来的是性能不再看运气,工程上这笔账很划算。

2. 节点存上高度

要判断树平不平衡,就得随时知道每个节点左右子树有多高。所以 AVL 树的节点比普通搜索树多存一个高度字段。代码里 BinaryNode 除了存值和左右孩子,还存 mHeight 和一个指向父节点的指针,旋转时改指针要靠父指针找到上头。

平衡因子就是某个节点左子树高度减右子树高度。按规矩,每个节点的平衡因子只能是 -1、0、1 这三个值,一旦变成 2 或者 -2,这个节点就是最小不平衡子树的根,必须在这里旋转。类的整体声明如下,后面几节逐个实现。

#pragma once

//二叉树节点
template<class T>
struct BinaryNode
{
	T mValue;
	int mHeight; //树的高度
	BinaryNode<T>* pLchild; //指向左子树节点
	BinaryNode<T>* pRchild; //指向右子树节点
	BinaryNode<T>* pParent; //指向父节点
};


template<class T>
class BinarySearchTree
{
public:
	// 构造函数
	BinarySearchTree();
	// 析构函数
	~BinarySearchTree();
	// 拷贝构造函数
	BinarySearchTree(const BinarySearchTree<T>& bst);

	//插入
	void Insert(const T& val);
	//查找
	BinaryNode<T>* SearchNode(const T& val);

	// 中序遍历
	template<class _Func>
	void Foreach(_Func func);
	// 辅助中序遍历
	template<class _Func>
	void MyForeach(BinaryNode<T>* root, _Func func);

	//释放二叉树
	void MyFree(BinaryNode<T>* root);
	//拷贝二叉树
	BinaryNode<T>* MyCopy(BinaryNode<T>* root, BinaryNode<T>* parent = NULL);

	//删除算法
	void Remove(const T& val);
	//第一种情况,如果删除的结点是叶子节点
	void RemoveNoChildNode(BinaryNode<T>* delNode);
	//第二种情况,如果删除的结点只有一个孩子
	void RemoveSingleChildNode(BinaryNode<T>* delNode);
	//第三种情况,如果删除的结点有两个孩子
	void RemoveDoubleChildNode(BinaryNode<T>* delNode);

	//获得最大值结点
	BinaryNode<T>* GetMaxNode(BinaryNode<T>* root);
	//获得最小值结点
	BinaryNode<T>* GetMinNode(BinaryNode<T>* root);
	/* 平衡二叉树接口 */
	//更新树的高度
	int UpdateTreeHeight(BinaryNode<T>* root);
	//获得某个子树的高度
	int GetTreeHeight(BinaryNode<T>* node);
	//求某个子树的平衡因子
	int GetBalanceFactor(BinaryNode<T>* node);
	// 右旋
	void RightRotate(BinaryNode<T>* node);
	// 左旋
	void leftRotate(BinaryNode<T>* node);
	// 调整树的平衡(subRoot)
	void AdjustTreeBalance(BinaryNode<T>* subRoot);

public:
	BinaryNode<T>* pRoot; //永远指向二叉排序树的根节点
	int mSize;
};

3. 四种旋转情形

发现某个节点平衡因子超了,就要旋转。失衡只可能出现在四种情况里,搞清楚是哪一种,对着选对应的转法就行。我们拿插入序列 3、1、2 实际走一遍,最能看清 LR 这种情况怎么来的。

我们先插 3,树根就是 3,高度 1。再插 1,1 比 3 小,挂在 3 的左孩子。此时 3 的左子树高 1、右子树高 0,平衡因子是 1,还在允许范围内。最后插 2,2 比 3 小走到左边,又比 1 大,挂在 1 的右孩子。这下 1 的右子树高 1、左子树高 0,1 的平衡因子变成 -1。而 3 的左子树现在高 2、右子树高 0,3 的平衡因子变成 2,超过 1 了,3 就是那个最小不平衡节点。

3 左边重,但它的左孩子 1 却是右边重,这种左孩子朝右歪的情况叫 LR。光把 3 右旋一下救不回来,得先把 1 左旋一次,让 2 顶上来,再把 3 右旋一次。转完之后根变成 2,2 的左孩子是 1,右孩子是 3,三个节点左右匀称,平衡因子全部回到 0。四种情形的判断和转法归拢成下表。

情形失衡点平衡因子较高一侧孩子的平衡因子旋转动作
LL 左左大于 1,左边重左孩子大于 0对失衡点右旋一次
LR 左右大于 1,左边重左孩子小于 0先左旋左孩子,再右旋失衡点
RR 右右小于 -1,右边重右孩子小于 0对失衡点左旋一次
RL 右左小于 -1,右边重右孩子大于 0先右旋右孩子,再左旋失衡点

代码里的 RightRotate 和 leftRotate 就是手动改指针。右旋时,把当前节点的左孩子托上来当新根,原来左孩子的右孙子过继给当前节点当左孩子,再把新旧父子关系接上。左旋是镜像操作。AdjustTreeBalance 从新插入的节点一路往上找父节点,第一个平衡因子绝对值超过 1 的就是失衡点,再按上表判断是哪种情形,调用对应的旋转组合。完整代码如下。

//树的旋转-向右旋转
template<class T>
void BinarySearchTree<T>::RightRotate(BinaryNode<T>* node)
{

	if (NULL == node)
	{
		return;
	}

	//缓存下当前子树的根节点的父节点
	BinaryNode<T>* pParent = node->pParent;
	//缓存下当前子树的根节点的左子树
	BinaryNode<T>* pLeftNode = node->pLchild;

	node->pLchild = pLeftNode->pRchild;
	if (pLeftNode->pRchild != NULL)
	{
		pLeftNode->pRchild->pParent = node;
	}


	pLeftNode->pRchild = node;
	node->pParent = pLeftNode;

	//备注:设置父节点
	pLeftNode->pParent = pParent;

	//父节点为NULL
	if (pParent == NULL)
	{
		pRoot = pLeftNode;
	}
	else
	{
		if (pParent->pLchild == node)
		{
			pParent->pLchild = pLeftNode;
		}
		else
		{
			pParent->pRchild = pLeftNode;
		}
	}

}


//左旋
template<class T>
void BinarySearchTree<T>::leftRotate(BinaryNode<T>* node)
{

	if (NULL == node)
	{
		return;
	}

	//缓存下父节点
	BinaryNode<T>* pParent = node->pParent;
	//缓存下右子树
	BinaryNode<T>* pRight = node->pRchild;

	node->pRchild = pRight->pLchild;
	if (pRight->pLchild != NULL)
	{
		pRight->pLchild->pParent = node;
	}


	pRight->pLchild = node;
	node->pParent = pRight;

	//备注:设置父节点
	pRight->pParent = pParent;

	if (pParent == NULL)
	{
		pRoot = pRight;
	}
	else
	{
		if (pParent->pLchild == node)
		{
			pParent->pLchild = pRight;
		}
		else
		{
			pParent->pRchild = pRight;
		}
	}

}

//调整树的平衡(subRoot)
template<class T>
void BinarySearchTree<T>::AdjustTreeBalance(BinaryNode<T>* subRoot)
{
	if (NULL == subRoot)
	{
		return;
	}

	//找到最小不平衡子树
	BinaryNode<T>* pBack = subRoot;
	//当前结点的平衡因子
	int factor = 0;
	while (pBack != NULL)
	{
		//计算当前结点的平衡因子
		factor = GetBalanceFactor(pBack);
		if (abs(factor) > 1)
		{
			break;
		}

		pBack = pBack->pParent;
	}


	if (pBack == NULL)
	{
		//此时树不需要调整
		return;
	}

	if (factor > 1)
	{
		//右旋
		if (GetBalanceFactor(pBack->pLchild) > 0)
		{
			//进行一次向右旋转
			RightRotate(pBack);
		}
		else
		{
			//先把pBack->pLchild左旋一次
			leftRotate(pBack->pLchild);
			//再把pBack整体向右旋转一次
			RightRotate(pBack);
		}
	}
	else
	{
		//左旋
		if (GetBalanceFactor(pBack->pRchild) < 0)
		{
			//直接向左旋转一次
			leftRotate(pBack);
		}
		else
		{
			//先把pBack->pRcild向右旋转一次
			RightRotate(pBack->pRchild);
			//然后再把pBack整体向左旋转一次
			leftRotate(pBack);
		}
	}
}

4. 插入后调平衡

旋转写好了,插入的流程和普通搜索树几乎一样,只是末尾多两步。先沿着树找到该挂的位置,把新节点接上,然后从根开始重新算一遍每个节点的高度,最后从新节点的父节点往上找不平衡点,交给 AdjustTreeBalance 处理。高度不更新,后面算平衡因子就全错,所以这一步省不掉。

//插入
template<class T>
void BinarySearchTree<T>::Insert(const T& val)
{
	if (SearchNode(val))
	{
		cout << "节点存在,拒绝插入!" << endl;
	}

	//指向插入位置的父节点
	BinaryNode<T>* pInsertNode = NULL;
	//指向当前结点
	BinaryNode<T>* pCurrent = pRoot;

	while (pCurrent != NULL)
	{
		pInsertNode = pCurrent;

		if (val < pCurrent->mValue)
		{
			pCurrent = pCurrent->pLchild;
		}
		else
		{
			pCurrent = pCurrent->pRchild;
		}
	}

	//创建新节点
	BinaryNode<T>* newnode = new BinaryNode<T>;
	newnode->mValue = val;
	newnode->pLchild = NULL;
	newnode->pRchild = NULL;
	newnode->pParent = pInsertNode;

	//判断pInsertNode是否是NULL,如果为NULL,说明当前树是一个空树
	if (pInsertNode == NULL)
	{
		pRoot = newnode;
	}
	else
	{
		if (val < pInsertNode->mValue)
		{
			pInsertNode->pLchild = newnode;
		}
		else
		{
			pInsertNode->pRchild = newnode;
		}
	}

	++mSize;

	//更新树的高度
	UpdateTreeHeight(pRoot);
	//调整树的平衡
	AdjustTreeBalance(newnode->pParent);
}

5. 删除的三情形

删除比插入麻烦,因为待删节点可能长在三种位置上,处理方式完全不同。第一种是叶子节点,左右孩子都空,直接把父节点指向它的指针断掉就行。第二种是只有一个孩子,把这个孩子提上去顶替待删节点的位置,接上父子关系。第三种是有两个孩子,这时候不能直接删,要找一个替补。通常取待删节点左子树里最大的那个节点,它一定没有右孩子,把它顶上来,再把它原来的位置处理掉。三种情况的代码如下。

//删除算法
template<class T>
void BinarySearchTree<T>::Remove(const T& val)
{
	BinaryNode<T>* pDelNode = SearchNode(val);

	if (pDelNode == NULL)
	{
		cout << "删除的结点不存在!" << endl;
		return;
	}

	//第一种情况,删除节点是叶子节点
	if (pDelNode->pLchild == NULL && pDelNode->pRchild == NULL)
	{
		RemoveNoChildNode(pDelNode);
	}
	//第三种情况,删除节点有两个孩子
	else if (pDelNode->pLchild != NULL && pDelNode->pRchild != NULL)
	{
		RemoveDoubleChildNode(pDelNode);
	}
	//第二种情况
	else
	{
		RemoveSingleChildNode(pDelNode);
	}

	//释放删除节点的内存
	delete pDelNode;
	pDelNode = NULL;
	--mSize;

}

//第一种情况,如果删除的结点是叶子节点
template<class T>
void BinarySearchTree<T>::RemoveNoChildNode(BinaryNode<T>* delNode)
{
	BinaryNode<T>* pDelNode = delNode;

	if (NULL == pDelNode)
	{
		return;
	}

	//判断删除的结点是否是根节点
	if (pDelNode == pRoot)
	{
		pRoot = NULL;
	}
	else
	{
		if (pDelNode->pParent->pLchild == pDelNode)
		{
			pDelNode->pParent->pLchild = NULL;
		}
		else
		{
			pDelNode->pParent->pRchild = NULL;
		}
	}
}

//第二种情况,如果删除的结点只有一个孩子
template<class T>
void BinarySearchTree<T>::RemoveSingleChildNode(BinaryNode<T>* delNode)
{
	BinaryNode<T>* pDelNode = delNode;

	if (NULL == pDelNode)
	{
		return;
	}

	if (pRoot == pDelNode)
	{
		//拿到删除节点的子节点
		BinaryNode<T>* pChild = (pRoot->pLchild == NULL ? pRoot->pRchild : pRoot->pLchild);
		//设置子节点的父节点为NULL
		pChild->pParent = NULL;
		//pRoot指向子节点
		pRoot = pChild;

	}
	else
	{
		BinaryNode<T>* pChild = (pDelNode->pLchild == NULL ? pDelNode->pRchild : pDelNode->pLchild);

		//更新pChild的父节点
		pChild->pParent = pDelNode->pParent;

		if (pDelNode->pParent->pLchild == pDelNode)
		{
			pDelNode->pParent->pLchild = pChild;
		}
		else
		{
			pDelNode->pParent->pRchild = pChild;
		}

	}

}
//第三种情况,如果删除的结点有两个孩子
template<class T>
void BinarySearchTree<T>::RemoveDoubleChildNode(BinaryNode<T>* delNode)
{
	BinaryNode<T>* pDelNode = delNode;

	if (NULL == pDelNode)
	{
		return;
	}

	//删除的结点是根节点的情况
	//pDelNode左子树最大值的结点,或者右子树中最小值的结点

	//获得删除节点左子树中最大的结点
	BinaryNode<T>* pLeftMaxNode = GetMaxNode(pDelNode->pLchild);
	if (pDelNode == pRoot)
	{

		if (pLeftMaxNode == pDelNode->pLchild)
		{
			pLeftMaxNode->pRchild = pRoot->pRchild;

		}
		else
		{
			//最大值的结点的父节点指向最大值结点的左子树
			pLeftMaxNode->pParent->pRchild = pLeftMaxNode->pLchild;
			if (pLeftMaxNode->pLchild != NULL)
			{
				pLeftMaxNode->pLchild->pParent = pLeftMaxNode->pParent;
			}
			//让最大值结点左右指针指向根节点的左右子树
			pLeftMaxNode->pLchild = pRoot->pLchild;
			pLeftMaxNode->pRchild = pRoot->pRchild;
		}

		pRoot = pLeftMaxNode;
	}
	else
	{

		if (pLeftMaxNode == pDelNode->pLchild)
		{
			pLeftMaxNode->pRchild = pDelNode->pRchild;
			//更新待删除结点的右子树的父节点
			pDelNode->pRchild->pParent = pLeftMaxNode;

			if (pDelNode->pParent->pLchild == pDelNode)
			{
				pDelNode->pParent->pLchild = pLeftMaxNode;
			}
			else
			{
				pDelNode->pParent->pRchild = pLeftMaxNode;
			}

		}
		else
		{
			pLeftMaxNode->pParent->pRchild = pLeftMaxNode->pLchild;
			if (pLeftMaxNode->pLchild != NULL)
			{
				pLeftMaxNode->pLchild->pParent = pLeftMaxNode->pParent;
			}

			pLeftMaxNode->pLchild = pDelNode->pLchild;
			pLeftMaxNode->pRchild = pDelNode->pRchild;

			if (pDelNode->pParent->pLchild == pDelNode)
			{
				pDelNode->pParent->pLchild = pLeftMaxNode;
			}
			else
			{
				pDelNode->pParent->pRchild = pLeftMaxNode;
			}
		}

	}
}

6. 配套基础接口

调平衡的核心讲完了,剩下几个接口是搭架子用的,和普通搜索树写法一样。构造函数把根置空,析构函数递归释放整棵树,中序遍历按左根右走,查找顺着比较结果一路往下。高度函数只是把节点存的 mHeight 读出来,空节点高度记为 0,平衡因子就是左右子树高度相减。这几块拼在一起,整棵 AVL 树就能跑了。

// 构造函数
template<class T>
BinarySearchTree<T>::BinarySearchTree()
{
	pRoot = NULL;
	mSize = 0;
}

// 析构函数
template<class T>
BinarySearchTree<T>::~BinarySearchTree()
{
	MyFree(pRoot);
}

// 拷贝构造
template<class T>
BinarySearchTree<T>::BinarySearchTree(const BinarySearchTree<T>& bst)
{
	mSize = bst.mSize;
	pRoot = MyCopy(bst.pRoot, NULL);
}

//释放二叉树
template<class T>
void BinarySearchTree<T>::MyFree(BinaryNode<T>* root)
{
	if (NULL == root)
	{
		return;
	}

	//释放左子树内存
	MyFree(root->pLchild);
	//释放右子树内存
	MyFree(root->pRchild);
	//释放当前结点的内存
	delete root;
}

//拷贝二叉树
template<class T>
BinaryNode<T>* BinarySearchTree<T>::MyCopy(BinaryNode<T>* root, BinaryNode<T>* parent = NULL)
{
	if (NULL == root)
	{
		return NULL;
	}

	//拷贝当前结点
	BinaryNode<T>* newnode = new BinaryNode<T>;
	newnode->mValue = root->mValue;
	newnode->pParent = parent;

	newnode->pLchild = NULL;
	newnode->pRchild = NULL;

	//拷贝左子树
	BinaryNode<T>* pLeft = MyCopy(root->pLchild, newnode);
	//拷贝右子树
	BinaryNode<T>* pRight = MyCopy(root->pRchild, newnode);

	//更新新节点的左右指针指向
	newnode->pLchild = pLeft;
	newnode->pRchild = pRight;

	return newnode;
}

// 中序遍历
template<class T> template<class _Func>
void BinarySearchTree<T>::Foreach(_Func func)
{
	MyForeach(pRoot, func);
}

//辅助中序遍历
template<class T> template<class _Func>
void BinarySearchTree<T>::MyForeach(BinaryNode<T>* root, _Func func)
{
	if (NULL == root)
	{
		return;
	}
	MyForeach(root->pLchild, func);
	func(root->mValue);
	MyForeach(root->pRchild, func);
}

//查找
template<class T>
BinaryNode<T>* BinarySearchTree<T>::SearchNode(const T& val)
{
	//辅助指针变量
	BinaryNode<T>* pCurrent = pRoot;
	while (pCurrent != NULL)
	{
		if (val < pCurrent->mValue)
		{
			pCurrent = pCurrent->pLchild;
		}
		else if (!(val < pCurrent->mValue) && !(pCurrent->mValue < val))
		{
			return pCurrent;
		}
		else
		{
			pCurrent = pCurrent->pRchild;
		}
	}

	return NULL;
}

//获得子树的高度
template<class T>
int BinarySearchTree<T>::GetTreeHeight(BinaryNode<T>* node)
{
	if (NULL == node)
	{
		return 0;
	}

	return node->mHeight;
}

//求某个子树的平衡因子
template<class T>
int BinarySearchTree<T>::GetBalanceFactor(BinaryNode<T>* node)
{
	if (NULL == node)
	{
		return 0;
	}

	return GetTreeHeight(node->pLchild) - GetTreeHeight(node->pRchild);
}

7. 常见理解误区

第一个误区,是把平衡因子的方向搞反。这里定义是左子树高度减右子树高度,正数代表左边重,负数代表右边重。判断 LL 还是 LR,必须拿失衡点孩子的符号再比一次,只看失衡点自己的符号会转反方向。

第二个误区,是以为旋转一次就够。LR 和 RL 这两种是先把歪的那一侧孩子转一下,再转失衡点,一共两次旋转。上来只转一次,树能好看一点但仍然不平衡,下一次插入还会再歪。

第三个误区,是以为 AVL 树把复杂度压到 O(log n) 就万事大吉。它确实把最坏情况从 O(n) 救回了 O(log n),但每次插入删除都要更新高度、往上找不平衡点,常数比普通搜索树大。查询密集、插入删除少的场景它最划算,频繁改动的场景后面还有更轻的平衡结构可选。

回到开头的问题。普通搜索树按序插入会退化成链表,查找掉到 O(n)。AVL 树靠节点存高度,用平衡因子盯住左右子树高度差不超过 1,一旦失衡就按 LL、LR、RR、RL 四种情形选一次或两次旋转掰回来,把插入、查找、删除都稳定在 O(log n)。插入 3、1、2 那个例子走通一遍,LR 为什么要转两次就真正懂了。实际工程里查询多读改写的场景,AVL 是把性能稳住的经典选择。