红黑树C实现

1. // red-black tree
3. #include <stdio.h>
4. #include <stdlib.h>
5. #include <string.h>
6. #include <stdarg.h>
7. #include <time.h>
8. //////////////////////
9. // supplied by user //
10. //////////////////////
12. typedef int KeyType;            // type of key
14. typedef struct {                // value related to key
15.     int stuff;
16. } ValType;
18. // how to compare keys
19. #define compLT(a,b) (a < b)
20. #define compEQ(a,b) (a == b)
22. /////////////////////////////////////
23. // implementation independent code //
24. /////////////////////////////////////
26. typedef enum {
27.     RBT_STATUS_OK,
28.     RBT_STATUS_MEM_EXHAUSTED,
29.     RBT_STATUS_DUPLICATE_KEY,
30.     RBT_STATUS_KEY_NOT_FOUND
31. } RbtStatus;
33. typedef enum { BLACK, RED } nodeColor;
35. typedef struct NodeTag {
36.     struct NodeTag *left;       // left child
37.     struct NodeTag *right;      // right child
38.     struct NodeTag *parent;     // parent
39.     nodeColor color;            // node color (BLACK, RED)
40.     KeyType key;                // key used for searching
41.     ValType val;                // data related to key
42. } NodeType;
44. #define SENTINEL &sentinel      // all leafs are sentinels(哨兵)
45. static NodeType sentinel = { SENTINEL, SENTINEL, 0, BLACK, 0};
47. static NodeType *root = SENTINEL; // root of red-black tree
49. static void rotateLeft(NodeType *x) {
51.     // rotate node x to left
53.     NodeType *y = x->right;
55.     // establish x->right link
56.     x->right = y->left;
57.     if (y->left != SENTINEL) y->left->parent = x;
59.     // establish y->parent link
60.     if (y != SENTINEL) y->parent = x->parent;
61.     if (x->parent) {
62.         if (x == x->parent->left)
63.             x->parent->left = y;
64.         else
65.             x->parent->right = y;
66.     } else {
67.         root = y;
68.     }
70.     // link x and y
71.     y->left = x;
72.     if (x != SENTINEL) x->parent = y;
73. }
75. static void rotateRight(NodeType *x) {
77.     // rotate node x to right
79.     NodeType *y = x->left;
81.     // establish x->left link
82.     x->left = y->right;
83.     if (y->right != SENTINEL) y->right->parent = x;
85.     // establish y->parent link
86.     if (y != SENTINEL) y->parent = x->parent;
87.     if (x->parent) {
88.         if (x == x->parent->right)
89.             x->parent->right = y;
90.         else
91.             x->parent->left = y;
92.     } else {
93.         root = y;
94.     }
96.     // link x and y
97.     y->right = x;
98.     if (x != SENTINEL) x->parent = y;
99. }
101. static void insertFixup(NodeType *x) {
103.     // maintain red-black tree balance
104.     // after inserting node x
106.     // check red-black properties
107.     while (x != root && x->parent->color == RED) {
108.         // we have a violation
109.         if (x->parent == x->parent->parent->left) {
110.             NodeType *y = x->parent->parent->right;
111.             if (y->color == RED) {
113.                 // uncle is RED
114.                 x->parent->color = BLACK;
115.                 y->color = BLACK;
116.                 x->parent->parent->color = RED;
117.                 x = x->parent->parent;
118.             } else {
120.                 // uncle is BLACK
121.                 if (x == x->parent->right) {
122.                     // make x a left child
123.                     x = x->parent;
124.                     rotateLeft(x);
125.                 }
127.                 // recolor and rotate
128.                 x->parent->color = BLACK;
129.                 x->parent->parent->color = RED;
130.                 rotateRight(x->parent->parent);
131.             }
132.         } else {
134.             // mirror image of above code
135.             NodeType *y = x->parent->parent->left;
136.             if (y->color == RED) {
138.                 // uncle is RED
139.                 x->parent->color = BLACK;
140.                 y->color = BLACK;
141.                 x->parent->parent->color = RED;
142.                 x = x->parent->parent;
143.             } else {
145.                 // uncle is BLACK
146.                 if (x == x->parent->left) {
147.                     x = x->parent;
148.                     rotateRight(x);
149.                 }
150.                 x->parent->color = BLACK;
151.                 x->parent->parent->color = RED;
152.                 rotateLeft(x->parent->parent);
153.             }
154.         }
155.     }
156.     root->color = BLACK;
157. }
159. // insert new node (no duplicates allowed)
160. RbtStatus rbtInsert(KeyType key, ValType val) {
161.     NodeType *current, *parent, *x;
163.     // allocate node for data and insert in tree
165.     // find future parent
166.     current = root;
167.     parent = 0;
168.     while (current != SENTINEL) {
169.         if (compEQ(key, current->key)) 
170.             return RBT_STATUS_DUPLICATE_KEY;
171.         parent = current;
172.         current = compLT(key, current->key) ?
173.             current->left : current->right;
174.     }
176.     // setup new node
177.     if ((x = malloc (sizeof(*x))) == 0)
178.         return RBT_STATUS_MEM_EXHAUSTED;
179.     x->parent = parent;
180.     x->left = SENTINEL;
181.     x->right = SENTINEL;
182.     x->color = RED;
183.     x->key = key;
184.     x->val = val;
186.     // insert node in tree
187.     if(parent) {
188.         if(compLT(key, parent->key))
189.             parent->left = x;
190.         else
191.             parent->right = x;
192.     } else {
193.         root = x;
194.     }
196.     insertFixup(x);
198.     return RBT_STATUS_OK;
199. }
201. static void deleteFixup(NodeType *x) {
203.     // maintain red-black tree balance
204.     // after deleting node x
206.     while (x != root && x->color == BLACK) {
207.         if (x == x->parent->left) {
208.             NodeType *w = x->parent->right;
209.             if (w->color == RED) {
210.                 w->color = BLACK;
211.                 x->parent->color = RED;
212.                 rotateLeft (x->parent);
213.                 w = x->parent->right;
214.             }
215.             if (w->left->color == BLACK && w->right->color == BLACK) {
216.                 w->color = RED;
217.                 x = x->parent;
218.             } else {
219.                 if (w->right->color == BLACK) {
220.                     w->left->color = BLACK;
221.                     w->color = RED;
222.                     rotateRight (w);
223.                     w = x->parent->right;
224.                 }
225.                 w->color = x->parent->color;
226.                 x->parent->color = BLACK;
227.                 w->right->color = BLACK;
228.                 rotateLeft (x->parent);
229.                 x = root;
230.             }
231.         } else {
232.             NodeType *w = x->parent->left;
233.             if (w->color == RED) {
234.                 w->color = BLACK;
235.                 x->parent->color = RED;
236.                 rotateRight (x->parent);
237.                 w = x->parent->left;
238.             }
239.             if (w->right->color == BLACK && w->left->color == BLACK) {
240.                 w->color = RED;
241.                 x = x->parent;
242.             } else {
243.                 if (w->left->color == BLACK) {
244.                     w->right->color = BLACK;
245.                     w->color = RED;
246.                     rotateLeft (w);
247.                     w = x->parent->left;
248.                 }
249.                 w->color = x->parent->color;
250.                 x->parent->color = BLACK;
251.                 w->left->color = BLACK;
252.                 rotateRight (x->parent);
253.                 x = root;
254.             }
255.         }
256.     }
257.     x->color = BLACK;
258. }
260. // delete node
261. RbtStatus rbtErase(NodeType * z) {
262.     NodeType *x, *y;
264.     if (z->left == SENTINEL || z->right == SENTINEL) {
265.         // y has a SENTINEL node as a child
266.         y = z;
267.     } else {
268.         // find tree successor with a SENTINEL node as a child
269.         y = z->right;
270.         while (y->left != SENTINEL) y = y->left;
271.     }
273.     // x is y's only child
274.     if (y->left != SENTINEL)
275.         x = y->left;
276.     else
277.         x = y->right;
279.     // remove y from the parent chain
280.     x->parent = y->parent;
281.     if (y->parent)
282.         if (y == y->parent->left)
283.             y->parent->left = x;
284.         else
285.             y->parent->right = x;
286.     else
287.         root = x;
289.     if (y != z) {
290.         z->key = y->key;
291.         z->val = y->val;
292.     }
295.     if (y->color == BLACK)
296.         deleteFixup (x);
298.     free (y);
300.     return RBT_STATUS_OK;
301. }
303. // find key
304. NodeType *rbtFind(KeyType key) {
305.     NodeType *current;
306.     current = root;
307.     while(current != SENTINEL) {
308.         if(compEQ(key, current->key)) {
309.             return current;
310.         } else {
311.             current = compLT (key, current->key) ?
312.                 current->left : current->right;
313.         }
314.     }
315.     return NULL;
316. }
318. // in-order walk of tree
319. void rbtInorder(NodeType *p, void (callback)(NodeType *)) {
320.     if (p == SENTINEL) return;
321.     rbtInorder(p->left, callback);
322.     callback(p);
323.     rbtInorder(p->right, callback);
324. }
326. // delete nodes depth-first
327. void rbtDelete(NodeType *p) {
328.     if (p == SENTINEL) return;
329.     rbtDelete(p->left);
330.     rbtDelete(p->right);
331.     free(p);
332. }
334. void displayNode(NodeType *p) {
335.     printf("%d %d/n", p->key, p->val.stuff);
336. }
338. int main(int argc, char **argv) {
339.     int maxnum, ct;
340.     KeyType key;
341.     RbtStatus status;
343.     // command-line:
344.     //
345.     //   rbt 2000
346.     //      process 2000 records
348.     NodeType *p;
350.     maxnum = atoi(argv[1]);
352.     printf("maxnum = %d/n", maxnum);
353.     for (ct = 1000; ct; ct--) {
354.         key = rand() % 1000 + 1;
355.         if ((p = rbtFind(key)) != NULL) {
356.    /*
357.                      if (p->val.stuff != 10*key) printf("fail val/n");
358.                         status = rbtErase(p);
359.                         if (status) printf("fail: status = %d/n", status);*/
360.             
361.         } else {
362.             ValType val;
363.             val.stuff = 10*key;
364.             status = rbtInsert(key, val);
365.             if (status) printf("fail: status = %d/n", status);
366.         }
367.     }
369.     for (ct = 1000; ct; ct--) {
370.         key = rand() % 100 + 1;
371.         if ((p = rbtFind(key)) != NULL) {
372.              status = rbtErase(p);
373.         }
374.     }
375.     // output nodes in order
376.     rbtInorder(root, displayNode);
378.     rbtDelete(root);
380.     return 0;
381. }