AVL树 模拟

AVL树 模拟

#include <utility>
#include <assert.h>

template <typename K, typename V>
struct AVLTreeNode
{
	AVLTreeNode(const pair<K, V>& kv)
		:_val(kv)
		,_left(nullptr)
		,_right(nullptr)
		,_parent(nullptr)
		,_bf(0)
	{}

	pair<K, V> _val;
	AVLTreeNode<K, V>* _left;
	AVLTreeNode<K, V>* _right;
	AVLTreeNode<K, V>* _parent;
	int _bf;
};

template <typename K, typename V>
class AVLTree
{
	typedef AVLTreeNode<K, V> Node;
public:
	AVLTree()
		:_root(nullptr)
	{}
	void RotateL(Node* parent)
	{
		Node* cur = parent->_right;
		Node* grandparent = parent->_parent;
		/*修改链接关系*/
		//grandparent
		if (parent == _root)
		{
			_root = cur;
		}
		else
		{
			if (grandparent->_left == parent)
			{
				grandparent->_left = cur;
			}
			else
			{
				grandparent->_right = cur;
			}
		}
		//cur左子树头节点
		if (cur->_left != nullptr)
		{
			cur->_left->_parent = parent;
		}
		//parent
		parent->_parent = cur;
		parent->_right = cur->_left;
		//cur
		cur->_left = parent;
		cur->_parent = grandparent;
		//平衡因子
		cur->_bf = parent->_bf = 0;
	}
	void RotateR(Node* parent)
	{
		Node* cur = parent->_left;
		Node* grandparent = parent->_parent;
		/*修改链接关系*/
		//grandparent
		if (parent == _root)
		{
			_root = cur;
		}
		else
		{
			if (grandparent->_left == parent)
			{
				grandparent->_left = cur;
			}
			else
			{
				grandparent->_right = cur;
			}
		}
		//cur右子树头节点
		if (cur->_right != nullptr)
		{
			cur->_right->_parent = parent;
		}
		//parent
		parent->_parent = cur;
		parent->_left = cur->_right;
		//cur
		cur->_parent = grandparent;
		cur->_right = parent;
		//平衡因子
		cur->_bf = parent->_bf = 0;
	}
	void RotateLR(Node* parent)
	{
		Node* cur = parent->_left;
		int bf = cur->_right->_bf;

		RotateL(cur);
		RotateR(parent);

		if (bf == 0)
		{
			parent->_bf = 0;
			cur->_bf = 0;
		}
		else if (bf == -1)
		{
			parent->_bf = 1;
			cur->_bf = 0;
		}
		else
		{
			parent->_bf = 0;
			cur->_bf = -1;
		}
	}
	void RotateRL(Node* parent)
	{
		Node* cur = parent->_right;
		int bf = cur->_left->_bf;

		RotateR(cur);
		RotateL(parent);

		if (bf == 0)
		{
			parent->_bf = 0;
			cur->_bf = 0;
		}
		else if (bf == -1)
		{
			parent->_bf = 0;
			cur->_bf = 1;
		}
		else
		{
			parent->_bf = -1;
			cur->_bf = 0;
		}
	}
	bool Insert(const pair<K, V>& kv)
	{
		if (_root == nullptr)
		{
			_root = new Node(kv);
			return true;
		}
		Node* cur = _root;
		Node* parent = nullptr;
		while (cur!=nullptr)
		{
			if (kv.first < (cur->_val).first)
			{
				parent = cur;
				cur = cur->_left;
			}
			else if (kv.first > (cur->_val).first)
			{
				parent = cur;
				cur = cur->_right;
			}
			else
			{
				return false;
			}
		}
		Node* newnode = new Node(kv);
		cur = newnode;
		if (kv.first < (parent->_val).first)
		{
			newnode->_parent = parent;
			parent->_left = newnode;
		}
		else if (kv.first > (parent->_val).first)
		{
			newnode->_parent = parent;
			parent->_right = newnode;
		}

		while (parent != nullptr)
		{
			//调整平衡因子
			if (parent->_left == cur)
			{
				parent->_bf--;
			}
			else if (parent->_right == cur)
			{
				parent->_bf++;
			}
			//根据调整后的平衡因子进行不同的决策
			if (parent->_bf == 0)
			{
				break;
			}
			else if (parent->_bf == 1 || parent->_bf == -1)
			{
				cur = parent;
				parent = parent->_parent;
			}
			else if (parent->_bf == 2 || parent->_bf == -2)
			{
				//旋转
				if (parent->_bf == 2 && cur->_bf == 1)
				{
					RotateL(parent);
				}
				else if (parent->_bf == 2 && cur->_bf == -1)
				{
					RotateRL(parent);
				}
				else if (parent->_bf == -2 && cur->_bf == 1)
				{
					RotateLR(parent);
				}
				else if (parent->_bf == -2 && cur->_bf == -1)
				{
					RotateR(parent);
				}
				break;
			}
			else
			{
				assert(false);
			}
		}
		return true;
	}
	bool IsAVLTree()
	{
		return _IsAVLTree(_root);
	}
private:
	int _Height(Node* root)
	{
		if (root == nullptr)
		{
			return 0;
		}
		int leftheight = _Height(root->_left);
		int rightheight = _Height(root->_right);
		return leftheight > rightheight ? leftheight + 1 : rightheight + 1;
	}
	bool _IsAVLTree(Node* root)
	{
		if (root == nullptr)
		{
			return true;
		}
		if (root->_left != nullptr && (root->_left->_val).first >= (root->_val).first)
		{
			return false;
		}
		if (root->_right != nullptr && (root->_right->_val).first <= (root->_val).first)
		{
			return false;
		}
		int leftheight = _Height(root->_left);
		int rightheight = _Height(root->_right);
		if (rightheight - leftheight != root->_bf)
		{
			return false;
		}
		return _IsAVLTree(root->_left) && _IsAVLTree(root->_right);
	}
	Node* _root;
};