git.iko.soy/ikoHSE/main
on master help sign in
ikoHSE/main/main.cpp
567 lines 21 kB history raw
1//
2// by Sergey Troshin, 27 March 2017.
3// With thankfulness to https://youtu.be/UaLIHuR1t8Q
4// But if smb would watch it
5// in DELETION video in case 5 parent can be both red and black, not just black.
6//
10#include <iostream>
11#include <vector>
12#include <set>
15namespace mystd {
16 template <typename T>
17 class set {
18 private:
19 enum color {
20 RED,
21 BLACK
22 };
23 struct node {
24 T val;
25 color cl = RED;
26 node * lch = nullptr;
27 node * rch = nullptr;
28 node * parent = nullptr;
29 bool is_root() {
30 return parent == nullptr;
31 }
32 node() {cl = RED;}
33 node(const T& val, color cl) : val(val), cl(cl) {}
34 node(const T& val, color cl, node* parent) : val(val), cl(cl), parent(parent) {}
35 node(T&& val, color cl) : val(val), cl(cl) {}
36 node(T&& val, color cl, node* parent) : val(val), cl(cl), parent(parent) {}
37 ~node() {
38 if (lch != nullptr) {
39 delete lch;
40 }
41 if (rch != nullptr) {
42 delete rch;
43 }
44 }
45 };
46 public:
47 class iterator;
48 node * root = nullptr;
49 node * sibling(node * me) {
50 if (me == nullptr || me->is_root()) return nullptr;
51 node * a = me->parent->lch;
52 node * b = me->parent->rch;
53 return (a == me ? b : a);
54 }
55 bool is_sibling_red(node * me) { // case 1
56 node * sib = sibling(me);
57 if (sib == nullptr) return 0;
58 return sib->cl == RED;
59 }
60 /**********************************************************************************
61 * if we find that we are red and our parent is red - *
62 * we should make some rotation in some cases *
63 * rotate(cur) basically suppose that cur is red and cur's parent is red *
64 **********************************************************************************/
65 void right_rotate(node * me, bool change_color) { // we are left child, so rotate our parent on positive angel
66 if (me->parent == root)
67 root = me;
68 node * parent = me->parent;
69 me->parent = parent->parent;
70 if (parent->parent != nullptr) {
71 if (parent->parent->lch == parent)
72 parent->parent->lch = me;
73 else
74 parent->parent->rch = me;
75 }
76 parent->lch = me->rch;
77 if (me->rch != nullptr)
78 me->rch->parent = parent;
79 me->rch = parent;
80 parent->parent = me;
81 if (change_color) {
82 parent->cl = RED;
83 me->cl = BLACK;
84 }
85 }
86 // same
87 void left_rotate(node * me, bool change_color) { // we are right child
88 if (me->parent == root)
89 root = me;
90 node * parent = me->parent;
91 me->parent = parent->parent; // me-to-grdad
92 if (parent->parent != nullptr) {
93 if (parent->parent->lch == parent)
94 parent->parent->lch = me; //grdad-to-me
95 else
96 parent->parent->rch = me;
97 }
98 parent->rch = me->lch;
99 if (me->lch != nullptr)
100 me->lch->parent = parent;
101 me->lch = parent;
102 parent->parent = me;
103 if (change_color) {
104 parent->cl = RED;
105 me->cl = BLACK;
106 }
107 }
108 //DEBUG ONLY
109 // calculate max black sum in subtree recursively
110 size_t black_sum(node * cur) {
111 if (cur == nullptr)
112 return 1;
113 size_t ans = std::max(black_sum(cur->lch), black_sum(cur->rch));
114 if (cur->cl == BLACK)
115 ans++;
116 return ans;
117 }
118 //DEBUG ONLY
119 // check red-black invariant recursively
120 bool is_tree_red_black(node * cur) {
121 if (cur == nullptr)
122 return 1;
123 if (root->cl == RED)
124 return 0;
125 // red-red relationship
126 if (get_color(cur) == RED && (get_color(cur->lch) == RED || get_color(cur->rch) == RED))
127 return 0;
128 return (black_sum(cur->lch) == black_sum(cur->rch)
129 && is_tree_red_black(cur->lch) && is_tree_red_black(cur->rch));
130 }
131 /********************************************************
132 * main function that repair invariant *
133 * it's fast and do constant number of rotations at all *
134 ********************************************************/
135 void check(node * cur) {
136 if (cur->parent == nullptr || cur->parent->cl == BLACK) {
137 return;
138 }
139 if (cur->parent->parent == nullptr) {
140 cur->parent->cl = BLACK;
141 return;
142 }
143 // so we are red, and our parent too
144 if (is_sibling_red(cur->parent)) {
145 node * sib = sibling(cur->parent);
146 cur->parent->cl = BLACK;
147 sib->cl = BLACK;
148 node * grand_dad = cur->parent->parent;
149 if (!grand_dad->is_root()) {
150 grand_dad->cl = RED;
151 cur = grand_dad;
152 check(cur);
153 } else {
154 return;
155 }
156 } else {
157 // now we face 4 cases
158 // LR means that we cur(red) is Left child and cur's parent(red too) is Right child
159 if (cur->parent->parent->lch == cur->parent && cur->parent->rch == cur) { // LR
160 left_rotate(cur, 0);
161 cur = cur->lch;
162 } else if (cur->parent->parent->rch == cur->parent && cur->parent->lch == cur) {// RL
163 right_rotate(cur, 0);
164 cur = cur->rch;
165 }
166 // only LL or RR ve have now. There we go.
167 if (cur->parent->parent->lch == cur->parent && cur->parent->lch == cur) { // LL
168 right_rotate(cur->parent, 1);
169 } else if (cur->parent->parent->rch == cur->parent && cur->parent->rch == cur) { // RR
170 left_rotate(cur->parent, 1);
171 } else {
172 throw std::invalid_argument("Undefined behaviour");
173 }
174 }
175 }
176 /****************************************
177 * INSERTION *
178 ****************************************/
179 void insert(T elem) {
180 node * new_node = new node(elem, RED);
181 if (root == nullptr) {
182 root = new_node;
183 root->cl = BLACK;
184 return;
185 }
186 node * curr = root;
187 /****************************************
188 * find a place to insert binary *
189 ****************************************/
190 while (1) {
191 if (new_node->val == curr->val) {
192 delete new_node;
193 return;
194 }
195 if (new_node->val < curr->val) { // идем либо влево, либо вправо
196 if (curr->lch == nullptr) {
197 curr->lch = new_node;
198 new_node->parent = curr;
199 break;
200 } else {
201 curr = curr->lch;
202 }
203 } else {
204 if (curr->rch == nullptr) {
205 curr->rch = new_node;
206 new_node->parent = curr;
207 break;
208 } else {
209 curr = curr->rch;
210 }
211 }
212 }
213 /****************************************
214 * then trying to repair tree invariant *
215 ****************************************/
216 // check if all is okey - parent is black
217 node * cur = new_node;
218 check(cur);
219 root->cl = BLACK;
220 }
221 size_t cnt_h(node * cur) {
222 if (cur == nullptr)
223 return 0;
224 return std::max(cnt_h(cur->lch), cnt_h(cur->rch)) + 1;
225 }
226 /********************************
227 * PRINT *
228 ********************************/
229 void print(node * cur, bool ind = 0) {
230 if (cur == nullptr)
231 return;
232 if (!ind)
233 std::cout << "\n---OUTPUT---\n";
234 std::cout << "col = " << (cur->cl == BLACK ? "black" : "red") <<
235 ", val = " << cur->val << "| [" << (cur->lch == nullptr ? -1 : cur->lch->val) <<
236 "," << (cur->rch == nullptr ? -1 : cur->rch->val) <<
237 "] parent = " << (cur->parent == nullptr ? -1 : cur->parent->val) << '\n';
238 print(cur->lch, 1);
239 print(cur->rch, 1);
240 }
241 ~set() {
242 if (root != nullptr) {
243 delete root;
244 }
245 }
246 bool empty() {
247 return root == nullptr;
248 }
249 node * find(node * cur, T& val) {
250 // if not found - return nullptr
251 // else return pointer on node with val
252 if (cur == nullptr)
253 return nullptr;
254 if (cur->val == val) {
255 return cur;
256 } else if (cur->val < val) {
257 return find(cur->rch, val);
258 } else {
259 return find(cur->lch, val);
260 }
261 }
262 node * last() {
263 node * r = root;
264 if (r == nullptr)
265 return r;
266 // find last element for end iterator
267 while (r->rch != nullptr)
268 r = r->rch;
269 return r;
270 }
271 iterator find(T& val) {
272 node * elem = find(root, val);
273 return elem == nullptr ? end() : iterator(elem, last());
274 }
275 node * child_to_delete(node * cur, bool is_left=0) {
276 // move 1 right and while can left
277 if (!is_left) {
278 if (cur->rch != nullptr)
279 return child_to_delete(cur->rch, 1);
280 else
281 return cur;
282 } else {
283 if (cur->lch == nullptr)
284 return cur;
285 else
286 return child_to_delete(cur->lch, 1);
287 }
288 }
289 void erase(T val) {
290 // if element exist remove it
291 node * elem = find(root, val);
292 if (elem == nullptr)
293 return;
294 node * me = child_to_delete(elem); // this node will be deleted implicitly
295 // copy value from the node which is about to be deleted...
296 // ...to the node which keep the val to be deleted
297 elem->val = me->val;
298 delete_node(me);
299 }
300 color get_color(node * n) {
301 if (n == nullptr)
302 return BLACK;
303 else return n->cl;
304 }
305 void set_color(node * n, color c) {
306 if (n == nullptr)
307 return;
308 else
309 n->cl = c;
310 }
311 bool is_left_child(node * sib) {
312 if (sib->parent == nullptr)
313 return 0;
314 return sib->parent->lch == sib;
315 }
316 /*******************************
317 * main is DELETE-function *
318 *******************************/
319 void delete_node(node * me) {
320 node * to_delete = me;
321 // if me is root
322 // * - means terminal state
323 if (me->parent == nullptr) {
324 if (me->lch != nullptr) {
325 root = me->lch;
326 root->parent = nullptr;
327 delete me; // we can have only one left right right child
328 } else if (me->rch != nullptr) {
329 root = me->rch;
330 root->parent = nullptr;
331 delete me;
332 } else {
333 root = nullptr;
334 }
335 return; // *
336 }
337 // if me is red node
338 if (me->cl == RED) {
339 // make me null double-black node
340 if (is_left_child(me))
341 me->parent->lch = nullptr;
342 else
343 me->parent->rch = nullptr;
344 delete to_delete;
345 return ; // *
346 }
347 // if me has red child
348 if (get_color(me->lch) == RED || get_color(me->rch) == RED) {
349 node * cop = nullptr;
350 if (get_color(me->lch) == RED) {
351 cop = me->lch;
352 } else {
353 cop = me->rch;
354 }
355 if (me->parent->rch == me) { // can delete me immediately
356 me->parent->rch = cop;
357 } else {
358 me->parent->lch = cop;
359 }
360 cop->parent = me->parent;
361 // make me null double-black node
362 if (is_left_child(to_delete))
363 to_delete->parent->lch = nullptr;
364 else
365 to_delete->parent->rch = nullptr;
366 delete to_delete;
367 return;
368 }
369 // if both children are black (it mean that they are null moreover)
370 // if we push the problem - we must solve it another time;
371 LOOP:;
372 // final state
373 if (me->parent == nullptr) {
374 if (is_left_child(to_delete))
375 to_delete->parent->lch = nullptr;
376 else
377 to_delete->parent->rch = nullptr;
378 delete to_delete;
379 return;
380 }
382 // case 1
383 // if our sibling is red, but parent is BLACK! - fix it and make some rotation
384 node * sib = sibling(me);
385 if (is_sibling_red(me) && get_color(me->parent) == BLACK && get_color(sib->lch) == BLACK && get_color(sib->rch) == BLACK) {
386 if (is_left_child(me)) {
387 left_rotate(sib, 1);
388 } else {
389 throw std::invalid_argument("Undefined behaviour");
390 }
391 }
393 //case 2
394 // if node's sibling is black and both sibling's children are black
395 // push double-black node up
396 sib = sibling(me);
397 if (get_color(me->parent) == BLACK && get_color(sib) == BLACK && get_color(sib->lch) == BLACK && get_color(sib->rch) == BLACK) {
398 // then we make pushing: sibling gonna be red, parent gonna be double red, me gonna be black
399 set_color(me, BLACK);
400 set_color(sib, RED);
401 set_color(me->parent, BLACK);
402 me = me->parent;
403 goto LOOP;
404 }
406 // case 3 - final state when our sibling and it's children are black but me's parent is red
407 if (!is_sibling_red(me) && get_color(me->parent) == RED && get_color(sib->lch) == BLACK && get_color(sib->rch) == BLACK) {
408 sib = sibling(me);
409 set_color(sib, RED);
410 set_color(me, BLACK);
411 set_color(me->parent, BLACK);
412 if (is_left_child(to_delete))
413 to_delete->parent->lch = nullptr;
414 else
415 to_delete->parent->rch = nullptr;
416 delete to_delete;
417 return;
418 }
420 //case 4
421 // if me is lch, and me's black sibling has left-red and right-black child
422 // we make right rotation from left-red child
423 // simmetric case if we are right child
424 sib = sibling(me);
425 if (get_color(sib) == BLACK) {
426 if (is_left_child(me) && get_color(sib->lch) == RED && get_color(sib->rch) == BLACK){ // usual case
427 right_rotate(sib->lch, 1);
428 } else if (!is_left_child(me) && get_color(sib->rch) == RED && get_color(sib->lch) == BLACK){ // simmetric
429 left_rotate(sib->rch, 1);
430 }
431 }
433 // case 5
434 // if me is left child, me's sibling is black, sibling's right child is red;
435 // simmetric if me is right child
436 sib = sibling(me);
437 if (get_color(sib) == BLACK) {
438 if (is_left_child(me) && get_color(sib->rch) == RED){ // usual case
439 left_rotate(sib, 0);
440 set_color(sib, me->parent->cl);
441 set_color(me->parent, BLACK);
442 set_color(sib->rch, BLACK);
443 } else if (!is_left_child(me) && get_color(sib->lch) == RED){ // simmetric
444 right_rotate(sib, 0);
445 set_color(sib, me->parent->cl);
446 set_color(me->parent, BLACK);
447 set_color(sib->lch, BLACK);
448 }
449 // make me null double-black node
450 if (is_left_child(to_delete))
451 to_delete->parent->lch = nullptr;
452 else
453 to_delete->parent->rch = nullptr;
454 delete to_delete;
455 return;
456 // *
457 }
458 goto LOOP;
459 return;
460 }
461 class iterator {
462 node * ptr;
463 public:
464 // is_end using for end iterator
465 bool is_end = false;
466 node * last;
467 iterator(node * ptr, node * last) : ptr(ptr), last(last) {}
468 node * next(node * p, node * prev = nullptr) {
469 if (ptr == last)
470 is_end = 1;
471 if (is_end)
472 return nullptr;
473 if (p->rch != nullptr && p->rch != prev) {
474 p = p->rch;
475 return p;
476 }
477 if (p->parent != nullptr) {
478 return next(p->parent, p);
479 }
480 return nullptr;
481 }
482 node * prev(node * p, node * prev = nullptr) {
483 if (is_end) {
484 is_end = 0;
485 return last;
486 }
487 if (p->lch != nullptr && p->lch != prev) {
488 p = p->lch;
489 return p;
490 }
491 if (p->parent != nullptr) {
492 return next(p->parent, p);
493 }
494 return nullptr;
495 }
496 iterator& operator ++ (int) {
497 ptr = next(ptr);
498 return (*this);
499 }
500 iterator operator ++ () {
501 auto it = *this;
502 ptr = next(ptr);
503 return it;
504 }
505 iterator& operator -- (int) {
506 ptr = prev(ptr);
507 return (*this);
508 }
509 iterator operator -- () {
510 auto it = *this;
511 ptr = prev(ptr);
512 return it;
513 }
514 node* operator -> () {
515 return ptr;
516 }
517 node* operator -> () const {
518 return ptr;
519 }
520 T& operator * () {
521 return ptr->val;
522 }
523 T& operator * () const {
524 return ptr->val;
525 }
526 bool operator == (const iterator& other) const {
527 return ptr == other.ptr;
528 }
529 bool operator != (const iterator& other) const {
530 return ptr != other.ptr;
531 }
532 bool null() {
533 return ptr == nullptr;
534 }
535 };
536 iterator begin() {
537 iterator it(root, last());
538 // find mostly left node - it'll be begin
539 while(it->lch != nullptr)
540 it--;
541 return it;
542 }
543 iterator end() {
544 iterator it(nullptr, last());
545 it.is_end = 1;
546 while(!it.null() && it->lch != nullptr)
547 it++;
548 return it;
549 }
550 };
552int main() {
553 mystd::set<int> st;
554 std::set<int> s;
555 int a;
556 int i = 0, n = 2000000;
557 while (i++ < n) {
558 a = rand() % n;
559 //std::cout << a << ' ';
560 st.insert(a);
561 }
562 while (!st.empty()) {
563 //st.print(st.root);
564 st.erase(*st.begin());
565 }
566 return 0;