平衡树之RB-tree

  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 };