Networking /C++/Linux

  1 /*-----------------------------------------------------------
  2     RB-Tree的插入和刪除操作的實現算法
  3     參考資料:
  4     1) <<Introduction to algorithm>>
  5     2) [url]http://lxr.linux.no/linux/lib/rbtree.c[/url]
  6 
  7     作者：[url]http://www.shnenglu.com/converse/[/url]
  8     您可以自由的傳播，修改這份代碼，轉載處請注明原作者
  9 
 10     紅黑樹的幾個性質:
 11     1) 每個結點只有紅和黑兩種顏色
 12     2) 根結點是黑色的
 13     3)空節點是黑色的（紅黑樹中，根節點的parent以及所有葉節點lchild、rchild都不指向NULL，而是指向一個定義好的空節點）。 
 14     4) 如果一個結點是紅色的,那么它的左右兩個子結點的顏色是黑色的
 15     5) 對于每個結點而言,從這個結點到葉子結點的任何路徑上的黑色結點
 16     的數目相同
 17 -------------------------------------------------------------*/
 18  
 19 #include <stdio.h>
 20 #include <stdlib.h>
 21 #include <string.h>
 22 
 23 typedef int key_t;
 24 typedef int data_t;
 25 
 26 typedef enum color_t
 27 {
 28     RED = 0,
 29     BLACK = 1
 30 }color_t;
 31 
 32 typedef struct rb_node_t
 33 {
 34     struct rb_node_t *left, *right, *parent;
 35     key_t key;
 36     data_t data;
 37     color_t color;
 38 }rb_node_t;
 39 
 40 /* forward declaration */
 41 rb_node_t* rb_insert(key_t key, data_t data, rb_node_t* root);
 42 rb_node_t* rb_search(key_t key, rb_node_t* root);
 43 rb_node_t* rb_erase(key_t key, rb_node_t* root);
 44 
 45 int main()
 46 {
 47     int i, count = 900000;
 48     key_t key;
 49     rb_node_t* root = NULL, *node = NULL;
 50     
 51     srand(time(NULL));
 52     for (i = 1; i < count; ++i)
 53     {
 54         key = rand() % count;
 55         if ((root = rb_insert(key, i, root)))
 56         {
 57             printf("[i = %d] insert key %d success!\n", i, key);
 58         }
 59         else
 60         {
 61             printf("[i = %d] insert key %d error!\n", i, key);
 62             exit(-1);
 63         }
 64 
 65         if ((node = rb_search(key, root)))
 66         {
 67             printf("[i = %d] search key %d success!\n", i, key);
 68         }
 69         else
 70         {
 71             printf("[i = %d] search key %d error!\n", i, key);
 72             exit(-1);
 73         }
 74         if (!(i % 10))
 75         {
 76             if ((root = rb_erase(key, root)))
 77             {
 78                 printf("[i = %d] erase key %d success\n", i, key);
 79             }
 80             else
 81             {
 82                 printf("[i = %d] erase key %d error\n", i, key);
 83             }
 84         }
 85     }
 86 
 87     return 0;
 88 }
 89 
 90 static rb_node_t* rb_new_node(key_t key, data_t data)
 91 {
 92     rb_node_t *node = (rb_node_t*)malloc(sizeof(struct rb_node_t));
 93 
 94     if (!node)
 95     {
 96         printf("malloc error!\n");
 97         exit(-1);
 98     }
 99     node->key = key, node->data = data;
100 
101     return node;
102 }
103 
104 /*-----------------------------------------------------------
105 |   node           right
106 |   / \    ==>     / \
107 |   a  right     node  y
108 |       / \           / \
109 |       b  y         a   b
110  -----------------------------------------------------------*/
111 static rb_node_t* rb_rotate_left(rb_node_t* node, rb_node_t* root)
112 {
113     rb_node_t* right = node->right;
114 
115     if ((node->right = right->left))
116     {
117         right->left->parent = node;
118     }
119     right->left = node;
120 
121     if ((right->parent = node->parent))
122     {
123         if (node == node->parent->right)
124         {
125             node->parent->right = right;
126         }
127         else
128         {
129             node->parent->left = right;
130         }
131     }
132     else
133     {
134         root = right;
135     }
136     node->parent = right;
137 
138     return root;
139 }
140 
141 /*-----------------------------------------------------------
142 |       node           left
143 |       / \            / \
144 |    left  y   ==>    a   node
145 |   / \               / \
146 |  a   b             b   y
147 -----------------------------------------------------------*/
148 static rb_node_t* rb_rotate_right(rb_node_t* node, rb_node_t* root)
149 {
150     rb_node_t* left = node->left;
151 
152     if ((node->left = left->right))
153     {
154         left->right->parent = node;
155     }
156     left->right = node;
157 
158     if ((left->parent = node->parent))
159     {
160         if (node == node->parent->right)
161         {
162             node->parent->right = left;
163         }
164         else
165         {
166             node->parent->left = left;
167         }
168     }
169     else
170     {
171         root = left;
172     }
173     node->parent = left;
174 
175     return root;
176 }
177 
178 static rb_node_t* rb_insert_rebalance(rb_node_t *node, rb_node_t *root)
179 {
180     rb_node_t *parent, *gparent, *uncle, *tmp;
181 
182     while ((parent = node->parent) && parent->color == RED)
183     {
184         gparent = parent->parent;
185 
186         if (parent == gparent->left)
187         {
188             uncle = gparent->right;
189             if (uncle && uncle->color == RED)
190             {
191                 uncle->color = BLACK;
192                 parent->color = BLACK;
193                 gparent->color = RED;
194                 node = gparent;
195             }
196             else
197             {
198                 if (parent->right == node)
199                 {
200                     root = rb_rotate_left(parent, root);
201                     tmp = parent;
202                     parent = node;
203                     node = tmp;
204                 }
205 
206                 parent->color = BLACK;
207                 gparent->color = RED;
208                 root = rb_rotate_right(gparent, root);
209             }
210         } 
211         else 
212         {
213             uncle = gparent->left;
214             if (uncle && uncle->color == RED)
215             {
216                 uncle->color = BLACK;
217                 parent->color = BLACK;
218                 gparent->color = RED;
219                 node = gparent;
220             }
221             else
222             {
223                 if (parent->left == node)
224                 {
225                     root = rb_rotate_right(parent, root);
226                     tmp = parent;
227                     parent = node;
228                     node = tmp;
229                 }
230 
231                 parent->color = BLACK;
232                 gparent->color = RED;
233                 root = rb_rotate_left(gparent, root);
234             }
235         }
236     }
237 
238     root->color = BLACK;
239 
240     return root;
241 }
242 
243 static rb_node_t* rb_erase_rebalance(rb_node_t *node, rb_node_t *parent, rb_node_t *root)
244 {
245     rb_node_t *other, *o_left, *o_right;
246 
247     while ((!node || node->color == BLACK) && node != root)
248     {
249         if (parent->left == node)
250         {
251             other = parent->right;
252             if (other->color == RED)
253             {
254                 other->color = BLACK;
255                 parent->color = RED;
256                 root = rb_rotate_left(parent, root);
257                 other = parent->right;
258             }
259             if ((!other->left || other->left->color == BLACK) &&
260                 (!other->right || other->right->color == BLACK))
261             {
262                 other->color = RED;
263                 node = parent;
264                 parent = node->parent;
265             }
266             else
267             {
268                 if (!other->right || other->right->color == BLACK)
269                 {
270                     if ((o_left = other->left))
271                     {
272                         o_left->color = BLACK;
273                     }
274                     other->color = RED;
275                     root = rb_rotate_right(other, root);
276                     other = parent->right;
277                 }
278                 other->color = parent->color;
279                 parent->color = BLACK;
280                 if (other->right)
281                 {
282                     other->right->color = BLACK;
283                 }
284                 root = rb_rotate_left(parent, root);
285                 node = root;
286                 break;
287             }
288         }
289         else
290         {
291             other = parent->left;
292             if (other->color == RED)
293             {
294                 other->color = BLACK;
295                 parent->color = RED;
296                 root = rb_rotate_right(parent, root);
297                 other = parent->left;
298             }
299             if ((!other->left || other->left->color == BLACK) &&
300                 (!other->right || other->right->color == BLACK))
301             {
302                 other->color = RED;
303                 node = parent;
304                 parent = node->parent;
305             }
306             else
307             {
308                 if (!other->left || other->left->color == BLACK)
309                 {
310                     if ((o_right = other->right))
311                     {
312                         o_right->color = BLACK;
313                     }
314                     other->color = RED;
315                     root = rb_rotate_left(other, root);
316                     other = parent->left;
317                 }
318                 other->color = parent->color;
319                 parent->color = BLACK;
320                 if (other->left)
321                 {
322                     other->left->color = BLACK;
323                 }
324                 root = rb_rotate_right(parent, root);
325                 node = root;
326                 break;
327             }
328         }
329     }
330 
331     if (node)
332     {
333         node->color = BLACK;
334     } 
335 
336     return root;
337 }
338 
339 static rb_node_t* rb_search_auxiliary(key_t key, rb_node_t* root, rb_node_t** save)
340 {
341     rb_node_t *node = root, *parent = NULL;
342     int ret;
343 
344     while (node)
345     {
346         parent = node;
347         ret = node->key - key;
348         if (0 < ret)
349         {
350             node = node->left;
351         }
352         else if (0 > ret)
353         {
354             node = node->right;
355         }
356         else
357         {
358             return node;
359         }
360     }
361 
362     if (save)
363     {
364         *save = parent;
365     }
366 
367     return NULL;
368 }
369 
370 rb_node_t* rb_insert(key_t key, data_t data, rb_node_t* root)
371 {
372     rb_node_t *parent = NULL, *node;
373 
374     parent = NULL;
375     if ((node = rb_search_auxiliary(key, root, &parent)))
376     {
377         return root;
378     }
379 
380     node = rb_new_node(key, data);
381     node->parent = parent; 
382     node->left = node->right = NULL;
383     node->color = RED;
384 
385     if (parent)
386     {
387         if (parent->key > key)
388         {
389             parent->left = node;
390         }
391         else
392         {
393             parent->right = node;
394         }
395     }
396     else
397     {
398         root = node;
399     }
400 
401     return rb_insert_rebalance(node, root);
402 }
403 
404 rb_node_t* rb_search(key_t key, rb_node_t* root)
405 {
406     return rb_search_auxiliary(key, root, NULL);
407 }
408 
409 rb_node_t* rb_erase(key_t key, rb_node_t *root)
410 {
411     rb_node_t *child, *parent, *old, *left, *node;
412     color_t color;
413 
414     if (!(node = rb_search_auxiliary(key, root, NULL)))
415     {
416         printf("key %d is not exist!\n");
417         return root;
418     }
419 
420     old = node;
421 
422     if (node->left && node->right)
423     {
424         node = node->right;
425         while ((left = node->left) != NULL)
426         {
427             node = left;
428         }
429         child = node->right;
430         parent = node->parent;
431         color = node->color;
432 
433         if (child)
434         {
435             child->parent = parent;
436         }
437         if (parent)
438         {
439             if (parent->left == node)
440             {
441                 parent->left = child;
442             }
443             else
444             {
445                 parent->right = child;
446             }
447         }
448         else
449         {
450             root = child;
451         }
452 
453         if (node->parent == old)
454         {
455             parent = node;
456         }
457 
458         node->parent = old->parent;
459         node->color = old->color;
460         node->right = old->right;
461         node->left = old->left;
462 
463         if (old->parent)
464         {
465             if (old->parent->left == old)
466             {
467                 old->parent->left = node;
468             }
469             else
470             {
471                 old->parent->right = node;
472             }
473         } 
474         else
475         {
476             root = node;
477         }
478 
479         old->left->parent = node;
480         if (old->right)
481         {
482             old->right->parent = node;
483         }
484     }
485     else
486     {
487         if (!node->left)
488         {
489             child = node->right;
490         }
491         else if (!node->right)
492         {
493             child = node->left;
494         }
495         parent = node->parent;
496         color = node->color;
497 
498         if (child)
499         {
500             child->parent = parent;
501         }
502         if (parent)
503         {
504             if (parent->left == node)
505             {
506                 parent->left = child;
507             }
508             else
509             {
510                 parent->right = child;
511             }
512         }
513         else
514         {
515             root = child;
516         }
517     }
518 
519     free(old);
520 
521     if (color == BLACK)
522     {
523         root = rb_erase_rebalance(child, parent, root);
524     }
525 
526     return root;
527 }