1 #include <memory>
2
3 template<class T>
4 struct rb_node
5 {
6 T key;
7 bool color;//true red | false black
8 std::shared_ptr<rb_node> lchild, rchild, parent;
9
10 rb_node(T key, bool color = true, std::shared_ptr<rb_node> lchild = nullptr,
11 std::shared_ptr<rb_node> rchild = nullptr, std::shared_ptr<rb_node> parent = nullptr)
12 :key(key)//此处不能使用this：现在对象还没被构建起来
13 {
14 //赋值，初始化在初始化列表中执行
15 this->color = color;
16 this->lchild = lchild;
17 this->rchild = rchild;
18 this->parent = parent;
19 }
20 };
21
22 template<class T>
23 class rb_tree
24 {
25 private:
26 std::shared_ptr<rb_node<T>> root;
27 std::shared_ptr<rb_node<T>> nil;//叶节点
28 void left_rotation(std::shared_ptr<rb_node<T>> a)
29 {
30 if (a == nullptr) return;
31 std::shared_ptr<rb_node<T>> b = a->rchild;
32 if (b == nullptr) return;
33 if (b->lchild != nullptr)
34 b->lchild->parent = a;
35 a->rchild = b->lchild;
36
37 if (a->parent == nil)
38 root = b;//rbtree的root以nil为父节点
39 else
40 {
41 if (a->parent->lchild == a)
42 a->parent->lchild = b;
43 else
44 a->parent->rchild = b;
45 }
46 b->parent = a->parent;
47
48 b->lchild = a;
49 a->parent = b;
50 }
51 void right_rotation(std::shared_ptr<rb_node<T>> a)
52 {
53 if (a == nullptr) return;
54 std::shared_ptr<rb_node<T>> b = a->lchild;
55 if (b == nullptr) return;
56 if (b->rchild != nullptr)
57 b->rchild->parent = a;
58 a->lchild = b->rchild;
59
60 if (a->parent == nil)
61 root = b;
62 else
63 {
64 if (a->parent->lchild == a)
65 a->parent->lchild = b;
66 else
67 a->parent->rchild = b;
68 }
69 b->parent = a->parent;
70
71 b->rchild = a;
72 a->parent = b;
73 }
74
75 public:
76 rb_tree()
77 {
78 root = nullptr;
79 T key;
80 nil = std::make_shared<rb_node<T>>(key, false);
81 }
82 void insert(T key)
83 {
84 std::shared_ptr<rb_node<T>> tmp = std::make_shared<rb_node<T>>(key, true, nil, nil, nil);
85 std::shared_ptr<rb_node<T>> ptr = root;
86
87 //情况1：树为空
88 if (ptr == nullptr)
89 {
90 tmp->color = false;
91 root = tmp;
92 return;
93 }
94 while (true)
95 {
96 if (key <= ptr->key)
97 {
98 if (ptr->lchild == nil) break;
99 ptr = ptr->lchild;
100 }
101 else
102 {
103 if (ptr->rchild == nil) break;
104 ptr = ptr->rchild;
105 }
106 }
107
108 if (key <= ptr->key)
109 ptr->lchild = tmp;
110 else
111 ptr->rchild = tmp;
112 tmp->parent = ptr;
113
114 while(true)
115 {
116 if (ptr == nil)//注意root可能被情况三修改为red，记得加特判
117 {
118 tmp->color = false;
119 root = tmp;
120 return;
121 }
122
123 //情况2：插入节点的父节点为黑色
124 if (!ptr->color) return;
125
126 /*情况3：插入节点的父节点和叔节点都存在且都为红色*/
127 if (ptr->parent->lchild->color && ptr->parent->rchild->color)
128 {
129 ptr->parent->color = true;
130 ptr->parent->lchild->color = false;
131 ptr->parent->rchild->color = false;
132
133 tmp = ptr->parent;
134 ptr = tmp->parent;
135 continue;
136 }
137 if (ptr->parent->lchild == ptr)
138 {
139 //情况4：右旋
140 if (tmp == ptr->lchild)
141 {
142 ptr->parent->color = true;
143 ptr->color = false;
144 right_rotation(ptr->parent);
145 return;
146 }
147 else
148 {
149 //情况5：左旋 + 右旋
150 left_rotation(ptr);
151 tmp = ptr;
152 ptr = tmp->parent;
153 continue;
154 }
155 }
156 else
157 {
158 //情况4：左旋
159 if (tmp == ptr->rchild)
160 {
161 ptr->parent->color = true;
162 ptr->color = false;
163 left_rotation(ptr->parent);
164 return;
165 }
166 else
167 {
168 //情况5：右旋+左旋
169 right_rotation(ptr);
170 tmp = ptr;
171 ptr = tmp->parent;
172 continue;
173 }
174 }
175 }
176 }
177 void earse(T key)
178 {
179 //待续
180 }
181 };
