STL05——手写一个简单版本的红黑树（500+行代码）

STL05——手写一个简单版本的红黑树

题目描述

在STL中，红黑树是一个重要的底层数据结构，本题需要设计一个 RedBlackTree 类，实现如下功能：

1、基础功能

• 构造函数：初始化 RedBlackTree 实例
• 析构函数：清理资源，确保无内存泄露

2、核心功能

• 在 RedBlackTree 中插入一个节点
• 在 RedBlackTree 中删除一个节点
• 查询 RedBlackTree 中某个节点是否存在
• 获取 RedBlackTree 中节点的数量
• 获取 RedBlackTree 中所有的节点

RedBlackTree 中的节点拥有两个属性，一个是 key，一个是 value，本题题意规定如果 key 相同，则代表是同一个节点。

输入描述

题目的包含多行输入，第一行为正整数 N, 代表后续有 N 行命令序列。

接下来 N 行，每行包含一个命令，命令格式为 [operation] [parameters] ，具体命令如下：

**
**

insert 命令：

• 格式：insert [key] [value]
• 功能：在 RedBlackTree 中添加节点，如果节点已经存在，则不进行任何操作

remove 命令：

• 格式：remove [key]
• 功能：删除 RedBlackTree 中的节点，如果节点不存在，则不进行任何操作

at 命令：

• 格式：at [key]
• 功能：查询 RedBlackTree 中的节点

size 命令：

• 格式：size
• 功能：获取 RedBlackTree 中节点的数量

print 命令：

• 格式：print
• 功能：按照中序遍历输出 RedBlackTree 中所有节点，格式为 [key1] [value1] [key2] [value2]

输出描述

输出为每行命令执行后的结果，具体输出格式如下：

insert 命令: 无输出

remove 命令: 无输出

at 命令: 输出一个整数，独占一行，代表 RedBlackTree 中 key 对应的 value，如果 key 不存在，则输出 not exsit

size 命令: 输出一个整数，独占一行，代表 RedBlackTree 中节点的数量

print 命令: 按照中序遍历输出 RedBlackTree 中所有节点的 key 和 value，格式为 [key1] [value1] [key2] [value2]…每个数字后都跟有一个空格，输出独占一行，如果 RedBlackTree 中不包含任何节点，则输出 empty

#include <iostream>
#include <sstream>
#include <string>
using namespace std;

enum class Color
{
	RED,
	BLACK
};

template<typename Key ,typename Value>
class RedBlackTree
{
	class Node 
	{
	public:
		Key key;
		Value value;
		Color color;
		Node* left;
		Node* right;
		Node* parent;
		
		Node()
			:color(Color::BLACK),left(NULL),right(NULL),parent(NULL)
		{}
		Node(const Key&key,const Value&value,Color color,Node*p=NULL)
			:key(key),value(value),color(color),left(NULL),right(NULL),parent(p)
		{}
	};

private:
	Node* root;
	size_t size;
	Node* Nil;//叶子结点

	//查询某结点
	Node* lookUp(const Key key)
	{
		Node* cmpNode = root;//从根节点开始
		while (cmpNode != NULL)
		{
			if (key < cmpNode->key)
				cmpNode = cmpNode->left;
			else if (key > cmpNode->key)
				cmpNode = cmpNode->right;
			else
				return cmpNode;
		}
		return NULL;
	}

	//右旋函数
	void rightRotate(Node* node)
	{
		Node* p = node->left;
		//处理p的右孩子
		node->left = p->right;
		if (p->right != NULL)
			p->right->parent = node;
		//把p置为老大
		p->parent = node->parent;
		if (node->parent == NULL)
			root = p;
		if (node == node->parent->left)
			node->parent->left = p;
		else if (node == node->parent->right)
			node->parent->right = p;
		//让node成为p的右孩子
		p->right = node;
		node->parent = p;
	}

	//左旋函数
	void leftRotate(Node* node)
	{
		Node* p = node->right;
		node->right = p->left;
		if (p->left != NULL)
			p->left->parent = node;
		p->parent = node->parent;
		if (node->parent == NULL)
			root = p;
		if (node == node->parent->left)
			node->parent->left = p;
		else if (node == node->parent->right)
			node->parent->right = p;
		p->left = node;
		node->parent = p;
	}

	//将插入的函数修复
	void insertFixup(Node*target)
	{
		while (target->parent != NULL && target->parent->color == Color::RED)
		{
			//父为祖左节点
			if (target->parent = target->parent->parent->left)
			{
				Node* uncle = target->parent->parent->right;
				//红叔
				if (uncle && uncle->color == Color::RED)
				{
					uncle->color == Color::BLACK;
					target->parent->color == Color::BLACK;
					target->parent->parent->color = Color::RED;
					target = target->parent->parent;
				}
				//黑叔
				else
				{
					//LR
					if (target == target->parent->right)
					{
						leftRotate(target->parent);
						rightRotate(target->parent);
						target->color = Color::BLACK;
						target->right->color = Color::RED;
					}
					//LL
					else if (target == target->parent->left)
					{
						rightRotate(target->parent->parent);
						target->parent->color = Color::BLACK;
						target->parent->right->color = Color::RED;
					}
				}
			}
			//父为祖的右节点
			else
			{
				Node* uncle = target->parent->left;
				//红叔
				if (uncle && uncle->color == Color::RED)
				{
					uncle->color == Color::BLACK;
					target->parent->color == Color::BLACK;
					target->parent->parent->color = Color::RED;
					target = target->parent->parent;
				}
				//黑叔
				else
				{
					//RL
					if (target == target->parent->left)
					{
						rightRotate(target->parent);
						leftRotate(target->parent);
						target->color = Color::BLACK;
						target->left->color = Color::RED;
					}
					//RR
					else if (target == target->parent->right)
					{
						leftRotate(target->parent->parent);
						target->parent->color = Color::BLACK;
						target->parent->left->color = Color::RED;
					}
				}	
			}
		}
		root->color = Color::BLACK;
	}
	void insertNode(const Key& key, const Value& value)
	{
		Node* newNode = new Node(key, value, Color::RED);
		Node* cur = root;
		if (root == NULL)
			root = newNode;
		while (cur->left!=NULL||cur->right!=NULL)
		{
			if (cur->key > newNode->key)
				cur = cur ->left;
			else if (cur->key < newNode->key)
				cur = cur->right;
			else
			{
				delete newNode;
				return;
			}
		}
		if (cur->key > newNode->key)
			cur->right = newNode;
		else
			cur->left = newNode;
		size++;
	}

	//中序遍历
	void inorderTraversal(Node* node) const
	{
		if (node != NULL)
		{
			inorderTraversal(node->left);
			cout << node->key << " ";
			cout << node->value << " ";
			inorderTraversal(node->right);
		}
	}
	
	//用新节点替换旧结点(旧结点不用再找了)
	void replaceNode(Node* target, Node* newNode)
	{
		if (target->parent == NULL)
			root = newNode;
		if (target == target->parent->left)
			target->parent->left = newNode;
		else if (target == target->parent->right)
			target->parent->right = newNode;
		if (newNode != NULL)
			newNode->parent = target->parent;
	}

	Node* findMinimumNode(Node* node)
	{
		while (node->left != NULL)
		{
			node = node->left;
		}
		return node;
	}

	// removeFixup函数用于在删除节点后恢复红黑树的性质
	void removeFixup(Node* node) {
		// 如果节点为Nil并且没有父节点，说明它是唯一的节点，直接返回
		if (node == Nil && node->parent == nullptr) {
			return;
		}

		// 当我们没有到达根节点时继续循环
		while (node != root) {
			// 如果节点是其父节点的左子节点
			if (node == node->parent->left) {
				// 兄弟节点是节点父亲的右子节点
				Node* sibling = node->parent->right;

				// 情况1：节点的兄弟节点是红色
				if (getColor(sibling) == Color::RED) {
					// 重新着色兄弟节点和父节点，并进行左旋
					setColor(sibling, Color::BLACK);
					setColor(node->parent, Color::RED);
					leftRotate(node->parent);
					// 旋转后更新兄弟节点
					sibling = node->parent->right;
				}

				// 情况2：兄弟节点的两个子节点都是黑色
				if (getColor(sibling->left) == Color::BLACK &&
					getColor(sibling->right) == Color::BLACK) {
					// 重新着色兄弟节点并向上移动
					setColor(sibling, Color::RED);
					node = node->parent;
					// 如果父节点是红色，将其改为黑色并结束
					if (node->color == Color::RED) {
						node->color = Color::BLACK;
						node = root;
					}
				}
				else {
					// 情况3：兄弟节点的右子节点是黑色（左子节点是红色）
					if (getColor(sibling->right) == Color::BLACK) {
						// 重新着色兄弟节点和兄弟节点的左子节点，并进行右旋
						setColor(sibling->left, Color::BLACK);
						setColor(sibling, Color::RED);
						rightRotate(sibling);
						// 旋转后更新兄弟节点
						sibling = node->parent->right;
					}

					// 情况4：兄弟节点的右子节点是红色
					setColor(sibling, getColor(node->parent));
					setColor(node->parent, Color::BLACK);
					setColor(sibling->right, Color::BLACK);
					leftRotate(node->parent);
					// 移动到根节点结束
					node = root;
				}
			}
			else {
				// 当节点是其父节点的右子节点时，对称的情况
				Node* sibling = node->parent->left;

				if (getColor(sibling) == Color::RED) {
					setColor(sibling, Color::BLACK);
					setColor(node->parent, Color::RED);
					rightRotate(node->parent);
					sibling = node->parent->left;
				}

				if (getColor(sibling->right) == Color::BLACK &&
					getColor(sibling->left) == Color::BLACK) {
					setColor(sibling, Color::RED);
					node = node->parent;
					if (node->color == Color::RED) {
						node->color = Color::BLACK;
						node = root;
					}
				}
				else {
					if (getColor(sibling->left) == Color::BLACK) {
						setColor(sibling->right, Color::BLACK);
						setColor(sibling, Color::RED);
						leftRotate(sibling);
						sibling = node->parent->left;
					}
					setColor(sibling, getColor(node->parent));
					setColor(node->parent, Color::BLACK);
					setColor(sibling->left, Color::BLACK);
					rightRotate(node->parent);
					node = root;
				}
			}
		}
		// 确保当前节点是黑色的，以维持红黑树性质
		setColor(node, Color::BLACK);
	}

	Color getColor(Node* node)
	{
		if (node == NULL)
			return Color::BLACK;
		else
			return node->color;
	}

	void setColor(Node* node, Color color) {
		if (node == nullptr) {
			return;
		}
		node->color = color;
	}
	//取消Nil的连接
	void dieConnectNil()
	{
		if (Nil == NULL)
			return;
		if (Nil->parent != NULL)
		{
			if (Nil == Nil->parent->left)
				Nil->parent->left = NULL;
			else
				Nil->parent->right = NULL;
		}
	}

	// 删除节点
	void deleteNode(Node* del) {
		Node* rep = del; // rep（替代节点）初始指向要删除的节点
		Node* child = nullptr;      // 要删除节点的孩子节点
		Node* parentRP;             // 替代节点的父节点
		Color origCol = rep->color; // 保存要删除节点的原始颜色

		// 如果删除节点没有左孩子
		if (!del->left) {
			rep = del->right;        // 替代节点指向删除节点的右孩子
			parentRP = del->parent;  // 更新替代节点的父节点
			origCol = getColor(rep); // 获取替代节点的颜色
			replaceNode(del, rep);   // 用替代节点替换删除节点
		}
		// 如果删除节点没有右孩子
		else if (!del->right) {
			rep = del->left;         // 替代节点指向删除节点的左孩子
			parentRP = del->parent;  // 更新替代节点的父节点
			origCol = getColor(rep); // 获取替代节点的颜色
			replaceNode(del, rep);   // 用替代节点替换删除节点
		}
		// 如果删除节点有两个孩子
		else {
			rep = findMinimumNode(
				del->right); // 找到删除节点右子树中的最小节点作为替代节点
			origCol = rep->color; // 保存替代节点的原始颜色
			// 如果替代节点不是删除节点的直接右孩子
			if (rep != del->right) {
				parentRP = rep->parent; // 更新替代节点的父节点
				child = rep->right; // 替代节点的右孩子变成要处理的孩子节点
				parentRP->left =
					child; // 替代节点的父节点的左孩子指向替代节点的孩子（因为替代节点是最小节点，所以不可能有左孩子）
				if (child != nullptr) {
					child->parent = parentRP; // 如果替代节点的孩子存在，则更新其父节点
				}
				// 将替代节点放到删除节点的位置
				del->left->parent = rep;
				del->right->parent = rep;
				rep->left = del->left;
				rep->right = del->right;
				// 如果删除节点有父节点，更新父节点的孩子指向
				if (del->parent != nullptr) {
					if (del == del->parent->left) {
						del->parent->left = rep;
						rep->parent = del->parent;
					}
					else {
						del->parent->right = rep;
						rep->parent = del->parent;
					}
				}
				// 如果删除节点没有父节点，说明它是根节点
				else {
					root = rep;
					root->parent = nullptr;
				}
			}
			// 如果替代节点是删除节点的直接右孩子
			else {
				child = rep->right; // 孩子节点指向替代节点的右孩子
				rep->left = del->left; // 替代节点的左孩子指向删除节点的左孩子
				del->left->parent = rep; // 更新左孩子的父节点
				// 更新删除节点父节点的孩子指向
				if (del->parent != nullptr) {
					if (del == del->parent->left) {
						del->parent->left = rep;
						rep->parent = del->parent;
					}
					else {
						del->parent->right = rep;
						rep->parent = del->parent;
					}
				}
				// 如果删除节点是根节点
				else {
					root = rep;
					root->parent = nullptr;
				}
				parentRP = rep; // 更新替代节点的父节点
			}
		}

		// 如果替代节点存在，更新其颜色为删除节点的颜色
		if (rep != nullptr) {
			rep->color = del->color;
		}
		// 如果替代节点不存在，将删除节点的颜色赋给origCol变量
		else {
			origCol = del->color;
		}

		// 如果原始颜色是黑色，需要进行额外的修复操作，因为黑色节点的删除可能会破坏红黑树的性质
		if (origCol == Color::BLACK) {
			// 如果存在孩子节点，进行修复操作
			if (child != nullptr) {
				removeFixup(child);
			}
			// 如果不存在孩子节点，将Nil节点（代表空节点）的父节点设置为替代节点的父节点
			else {
				Nil->parent = parentRP;
				// 如果替代节点的父节点存在，设置其对应的孩子指针为Nil节点
				if (parentRP != nullptr) {
					if (parentRP->left == nullptr) {
						parentRP->left = Nil;
					}
					else {
						parentRP->right = Nil;
					}
				}
				// 进行修复操作
				removeFixup(Nil);
				// 断开Nil节点与树的连接，因为在红黑树中Nil节点通常是单独存在的
				dieConnectNil();
			}
		}

		// 删除节点
		delete del;
	}

	public:
		RedBlackTree() :root(NULL), size(0), Nil(new Node())
		{
			Nil->color = Color::BLACK;
		}

		~RedBlackTree()
		{
			deleteTree(root);
		}

		void insert(const Key& key, const Value& value) 
		{ 
			insertNode(key, value); 
		}

		void remove(const Key& key)
		{
			Node* p = lookUp(key);
			if (p != NULL)
			{
				deleteNode(p);
				size--;
			}
		}

		Value* at(const Key& key)
		{
			auto ans = lookUp(key);
			if (ans == NULL)
				return NULL;
			else
				return &ans->value;
		}

		int getSize() { return size; }

		bool empty() { return size == 0; }

		void print() {
			inorderTraversal(root);
			std::cout << std::endl;
		}

		void clear()
		{
			deleteNode(root);
			size = 0;
		}
		private:
			void deleteTree(Node* node)
			{
				if (node)
				{
					deleteTree(node->left);
					deleteTree(node->right);
					delete node;
				}
			}
};
int main() {
	// 创建红黑树实例
	RedBlackTree<int, int> rbTree;

	int N;
	std::cin >> N;
	getchar();

	std::string line;
	for (int i = 0; i < N; i++)
	{
		std::getline(std::cin, line);
		std::istringstream iss(line);
		std::string command;
		iss >> command;

		int key;
		int value;

		if (command == "insert")
		{
			iss >> key >> value;
			rbTree.insert(key, value);
		}

		if (command == "size")
		{
			std::cout << rbTree.getSize() << std::endl;
		}

		if (command == "at")
		{
			iss >> key;
			int* res = rbTree.at(key);
			if (res == nullptr)
			{
				std::cout << "not exist" << std::endl;
			}
			else
			{
				std::cout << *res << std::endl;
			}
		}

		if (command == "remove")
		{
			iss >> key;
			rbTree.remove(key);
		}

		if (command == "print")
		{
			if (rbTree.empty())
			{
				std::cout << "empty" << std::endl;
			}
			else
			{
				rbTree.print();
			}
		}
	}
	return 0;
}