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;elsereturn 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{//LRif (target == target->parent->right){leftRotate(target->parent);rightRotate(target->parent);target->color = Color::BLACK;target->right->color = Color::RED;}//LLelse 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{//RLif (target == target->parent->left){rightRotate(target->parent);leftRotate(target->parent);target->color = Color::BLACK;target->left->color = Color::RED;}//RRelse 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;elsecur->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;elsereturn 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;elseNil->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;elsereturn &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;
}