平衡数之Treap

  1 #include <memory>//智能指针头文件
  2 #include <random>//随机数头文件
  3 #include <iostream>
  4 #include <algorithm>
  5 #include <ctime>//time头文件
  6 
  7 template<class T>
  8 struct node
  9 {
 10     T key;
 11     unsigned weight;
 12     std::shared_ptr<node<T>> lchild, rchild, parent;
 13     node(T key, unsigned weight, std::shared_ptr<node<T>> lchild = nullptr, 
 14         std::shared_ptr<node<T>> rchild = nullptr, std::shared_ptr<node<T>> parent = nullptr):
 15         key(key)
 16     {
 17         this->weight = weight;
 18         this->lchild = lchild;
 19         this->rchild = rchild;
 20         this->parent = parent;
 21     }
 22 };
 23 
 24 template<class T>
 25 class Treap
 26 {
 27 private:
 28     std::shared_ptr<node<T>> root;
 29     std::default_random_engine generator;
 30     std::uniform_int_distribution<int> dis;
 31     void left_rotation(std::shared_ptr<node<T>> a)
 32     {
 33         if (a == nullptr) return;
 34 
 35         std::shared_ptr<node<T>> b = a->lchild;
 36         if (b == nullptr) return;
 37 
 38         if (b->rchild != nullptr)
 39         {
 40             a->lchild = b->rchild;
 41             b->rchild->parent = a;
 42         }
 43         else
 44             a->lchild = nullptr;//
 45 
 46         b->parent = a->parent;
 47         if (a->parent == nullptr)
 48             root = b;
 49         else
 50         {
 51             if (a->parent->lchild == a)
 52                 a->parent->lchild = b;
 53             else
 54                 a->parent->rchild = b;
 55         }
 56 
 57         b->rchild = a;
 58         a->parent = b;
 59     }
 60     void right_rotation(std::shared_ptr<node<T>> a)
 61     {
 62         if (a == nullptr) return;
 63 
 64         std::shared_ptr<node<T>> b = a->rchild;
 65         if (b == nullptr) return;
 66 
 67         if (b->lchild != nullptr)
 68         {
 69             a->rchild = b->lchild;
 70             b->lchild->parent = a;
 71         }
 72         else
 73             a->rchild = nullptr;
 74 
 75         b->parent = a->parent;
 76         if (a->parent == nullptr)
 77             root = b;
 78         else
 79         {
 80             if (a->parent->lchild == a)
 81                 a->parent->lchild = b;
 82             else
 83                 a->parent->rchild = b;
 84         }
 85 
 86         b->lchild = a;
 87         a->parent = b;
 88     }
 89 public:
 90     Treap() :generator(time(nullptr)), dis(0, 100000)
 91     {
 92         root = nullptr;
 93     }
 94     void insert(T key)
 95     {
 96         std::shared_ptr<node<T>> tmp = std::make_shared<node<T>>(key, dis(generator));
 97         std::shared_ptr<node<T>> ptr = root;
 98         if (ptr == nullptr)
 99         {
100             root = tmp;
101             return;
102         }
103         
104         while (true)
105         {
106             if (key <= ptr->key)
107             {
108                 if (ptr->lchild == nullptr) break;
109                 ptr = ptr->lchild;
110             }
111             else
112             {
113                 if (ptr->rchild == nullptr) break;
114                 ptr = ptr->rchild;
115             }
116         }
117 
118         if (key <= ptr->key)
119             ptr->lchild = tmp;
120         else
121             ptr->rchild = tmp;
122         tmp->parent = ptr;
123 
124         while (ptr != nullptr)
125         {
126             if (tmp->weight < ptr->weight)
127             {
128                 if (ptr->lchild == tmp)
129                 {
130                     left_rotation(ptr);
131                     ptr = tmp->parent;
132                 }
133                 else
134                 {
135                     right_rotation(ptr);
136                     ptr = tmp->parent;
137                 }
138             }
139             else
140                 break;
141         }
142     }
143     void earse(T key)
144     {
145         std::shared_ptr<node<T>> ptr = root;
146         while (ptr != nullptr)
147         {
148             if (key == ptr->key)
149                 break;
150             else if (key < ptr->key)
151                 ptr = ptr->lchild;
152             else
153                 ptr = ptr->rchild;
154         }
155         if (ptr == nullptr) return;
156         while (true)
157         {
158             if (ptr->lchild == nullptr || ptr->rchild == nullptr)
159             {
160                 if (ptr->lchild == nullptr && ptr->rchild == nullptr)
161                 {
162                     if (ptr->parent == nullptr) root = nullptr;
163                     else
164                     {
165                         if (ptr->parent->lchild == ptr)
166                             ptr->parent->lchild = nullptr;
167                         else
168                             ptr->parent->rchild = nullptr;
169                     }
170                 }
171                 else if (ptr->lchild == nullptr)
172                 {
173                     if (ptr->parent == nullptr)
174                         root = ptr->rchild;
175                     else
176                     {
177                         if (ptr->parent->lchild == ptr)
178                             ptr->parent->lchild = ptr->rchild;
179                         else
180                             ptr->parent->rchild = ptr->rchild;
181                     }
182                     ptr->rchild->parent = ptr->parent;
183                 }
184                 else
185                 {
186                     if (ptr->parent == nullptr)
187                         root = ptr->lchild;
188                     else
189                     {
190                         if (ptr->parent->lchild == ptr)
191                             ptr->parent->lchild = ptr->lchild;
192                         else
193                             ptr->parent->rchild = ptr->lchild;
194                     }
195                     ptr->lchild->parent = ptr->parent;
196                 }
197                 return;
198             }
199             else
200             {
201                 if (ptr->lchild->weight <= ptr->rchild->weight)
202                     left_rotation(ptr);
203                 else
204                     right_rotation(ptr);
205             }
206         }
207     }
208 };
原文地址:https://www.cnblogs.com/txlstars/p/5655672.html