rbtree
view src/rbtree.c @ 14:1b8a3a6088b6
simplified find_min
author | John Tsiombikas <nuclear@member.fsf.org> |
---|---|
date | Sun, 02 Nov 2014 10:56:08 +0200 |
parents | 5b955b3a4675 |
children | 1b77b72688fe |
line source
1 /*
2 rbtree - simple balanced binary search tree (red-black tree) library.
3 Copyright (C) 2011-2014 John Tsiombikas <nuclear@member.fsf.org>
5 rbtree is free software, feel free to use, modify, and redistribute it, under
6 the terms of the 3-clause BSD license. See COPYING for details.
7 */
8 #include <stdio.h>
9 #include <stdlib.h>
10 #include <stdint.h>
11 #include <string.h>
12 #include "rbtree.h"
14 #define INT2PTR(x) ((void*)(intptr_t)(x))
15 #define PTR2INT(x) ((int)(intptr_t)(x))
17 struct rbtree {
18 struct rbnode *root;
20 rb_alloc_func_t alloc;
21 rb_free_func_t free;
23 rb_cmp_func_t cmp;
24 rb_del_func_t del;
25 void *del_cls;
27 struct rbnode *rstack, *iter;
28 };
30 static int cmpaddr(const void *ap, const void *bp);
31 static int cmpint(const void *ap, const void *bp);
33 static int count_nodes(struct rbnode *node);
34 static void del_tree(struct rbnode *node, void (*delfunc)(struct rbnode*, void*), void *cls);
35 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data);
36 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key);
37 /*static struct rbnode *find(struct rbtree *rb, struct rbnode *node, void *key);*/
38 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls);
40 struct rbtree *rb_create(rb_cmp_func_t cmp_func)
41 {
42 struct rbtree *rb;
44 if(!(rb = malloc(sizeof *rb))) {
45 return 0;
46 }
47 if(rb_init(rb, cmp_func) == -1) {
48 free(rb);
49 return 0;
50 }
51 return rb;
52 }
54 void rb_free(struct rbtree *rb)
55 {
56 rb_destroy(rb);
57 free(rb);
58 }
61 int rb_init(struct rbtree *rb, rb_cmp_func_t cmp_func)
62 {
63 memset(rb, 0, sizeof *rb);
65 if(cmp_func == RB_KEY_INT) {
66 rb->cmp = cmpint;
67 } else if(cmp_func == RB_KEY_STRING) {
68 rb->cmp = (rb_cmp_func_t)strcmp;
69 } else {
70 rb->cmp = cmpaddr;
71 }
73 rb->alloc = malloc;
74 rb->free = free;
75 return 0;
76 }
78 void rb_destroy(struct rbtree *rb)
79 {
80 del_tree(rb->root, rb->del, rb->del_cls);
81 }
83 void rb_set_allocator(struct rbtree *rb, rb_alloc_func_t alloc, rb_free_func_t free)
84 {
85 rb->alloc = alloc;
86 rb->free = free;
87 }
90 void rb_set_compare_func(struct rbtree *rb, rb_cmp_func_t func)
91 {
92 rb->cmp = func;
93 }
95 void rb_set_delete_func(struct rbtree *rb, rb_del_func_t func, void *cls)
96 {
97 rb->del = func;
98 rb->del_cls = cls;
99 }
102 void rb_clear(struct rbtree *rb)
103 {
104 del_tree(rb->root, rb->del, rb->del_cls);
105 rb->root = 0;
106 }
108 int rb_copy(struct rbtree *dest, struct rbtree *src)
109 {
110 struct rbnode *node;
112 rb_clear(dest);
113 rb_begin(src);
114 while((node = rb_next(src))) {
115 if(rb_insert(dest, node->key, node->data) == -1) {
116 return -1;
117 }
118 }
119 return 0;
120 }
122 int rb_size(struct rbtree *rb)
123 {
124 return count_nodes(rb->root);
125 }
127 int rb_insert(struct rbtree *rb, void *key, void *data)
128 {
129 rb->root = insert(rb, rb->root, key, data);
130 rb->root->red = 0;
131 return 0;
132 }
134 int rb_inserti(struct rbtree *rb, int key, void *data)
135 {
136 rb->root = insert(rb, rb->root, INT2PTR(key), data);
137 rb->root->red = 0;
138 return 0;
139 }
142 int rb_delete(struct rbtree *rb, void *key)
143 {
144 rb->root = delete(rb, rb->root, key);
145 rb->root->red = 0;
146 return 0;
147 }
149 int rb_deletei(struct rbtree *rb, int key)
150 {
151 rb->root = delete(rb, rb->root, INT2PTR(key));
152 rb->root->red = 0;
153 return 0;
154 }
157 void *rb_find(struct rbtree *rb, void *key)
158 {
159 struct rbnode *node = rb->root;
161 while(node) {
162 int cmp = rb->cmp(key, node->key);
163 if(cmp == 0) {
164 return node;
165 }
166 node = cmp < 0 ? node->left : node->right;
167 }
168 return 0;
169 }
171 void *rb_findi(struct rbtree *rb, int key)
172 {
173 return rb_find(rb, INT2PTR(key));
174 }
177 void rb_foreach(struct rbtree *rb, void (*func)(struct rbnode*, void*), void *cls)
178 {
179 traverse(rb->root, func, cls);
180 }
183 struct rbnode *rb_root(struct rbtree *rb)
184 {
185 return rb->root;
186 }
188 void rb_begin(struct rbtree *rb)
189 {
190 rb->rstack = 0;
191 rb->iter = rb->root;
192 }
194 #define push(sp, x) ((x)->next = (sp), (sp) = (x))
195 #define pop(sp) ((sp) = (sp)->next)
196 #define top(sp) (sp)
198 struct rbnode *rb_next(struct rbtree *rb)
199 {
200 struct rbnode *res = 0;
202 while(rb->rstack || rb->iter) {
203 if(rb->iter) {
204 push(rb->rstack, rb->iter);
205 rb->iter = rb->iter->left;
206 } else {
207 rb->iter = top(rb->rstack);
208 pop(rb->rstack);
209 res = rb->iter;
210 rb->iter = rb->iter->right;
211 break;
212 }
213 }
214 return res;
215 }
217 void *rb_node_key(struct rbnode *node)
218 {
219 return node ? node->key : 0;
220 }
222 int rb_node_keyi(struct rbnode *node)
223 {
224 return node ? PTR2INT(node->key) : 0;
225 }
227 void *rb_node_data(struct rbnode *node)
228 {
229 return node ? node->data : 0;
230 }
232 static int cmpaddr(const void *ap, const void *bp)
233 {
234 return ap < bp ? -1 : (ap > bp ? 1 : 0);
235 }
237 static int cmpint(const void *ap, const void *bp)
238 {
239 return PTR2INT(ap) - PTR2INT(bp);
240 }
243 /* ---- left-leaning 2-3 red-black implementation ---- */
245 /* helper prototypes */
246 static int is_red(struct rbnode *tree);
247 static void color_flip(struct rbnode *tree);
248 static struct rbnode *rot_left(struct rbnode *a);
249 static struct rbnode *rot_right(struct rbnode *a);
250 static struct rbnode *find_min(struct rbnode *tree);
251 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree);
252 /*static struct rbnode *move_red_right(struct rbnode *tree);*/
253 static struct rbnode *move_red_left(struct rbnode *tree);
254 static struct rbnode *fix_up(struct rbnode *tree);
256 static int count_nodes(struct rbnode *node)
257 {
258 if(!node)
259 return 0;
261 return 1 + count_nodes(node->left) + count_nodes(node->right);
262 }
264 static void del_tree(struct rbnode *node, rb_del_func_t delfunc, void *cls)
265 {
266 if(!node)
267 return;
269 del_tree(node->left, delfunc, cls);
270 del_tree(node->right, delfunc, cls);
272 if(delfunc) {
273 delfunc(node, cls);
274 }
275 free(node);
276 }
278 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data)
279 {
280 int cmp;
282 if(!tree) {
283 struct rbnode *node = rb->alloc(sizeof *node);
284 node->red = 1;
285 node->key = key;
286 node->data = data;
287 node->left = node->right = 0;
288 return node;
289 }
291 cmp = rb->cmp(key, tree->key);
293 if(cmp < 0) {
294 tree->left = insert(rb, tree->left, key, data);
295 } else if(cmp > 0) {
296 tree->right = insert(rb, tree->right, key, data);
297 } else {
298 tree->data = data;
299 }
301 /* fix right-leaning reds */
302 if(is_red(tree->right)) {
303 tree = rot_left(tree);
304 }
305 /* fix two reds in a row */
306 if(is_red(tree->left) && is_red(tree->left->left)) {
307 tree = rot_right(tree);
308 }
310 /* if 4-node, split it by color inversion */
311 if(is_red(tree->left) && is_red(tree->right)) {
312 color_flip(tree);
313 }
315 return tree;
316 }
318 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key)
319 {
320 int cmp;
322 if(!tree) {
323 return 0;
324 }
326 cmp = rb->cmp(key, tree->key);
328 if(cmp < 0) {
329 if(!is_red(tree->left) && !is_red(tree->left->left)) {
330 tree = move_red_left(tree);
331 }
332 tree->left = delete(rb, tree->left, key);
333 } else {
334 /* need reds on the right */
335 if(is_red(tree->left)) {
336 tree = rot_right(tree);
337 }
339 /* found it at the bottom (XXX what certifies left is null?) */
340 if(cmp == 0 && !tree->right) {
341 if(rb->del) {
342 rb->del(tree, rb->del_cls);
343 }
344 rb->free(tree);
345 return 0;
346 }
348 if(!is_red(tree->right) && !is_red(tree->right->left)) {
349 tree = move_red_left(tree);
350 }
352 if(key == tree->key) {
353 struct rbnode *rmin = find_min(tree->right);
354 tree->key = rmin->key;
355 tree->data = rmin->data;
356 tree->right = del_min(rb, tree->right);
357 } else {
358 tree->right = delete(rb, tree->right, key);
359 }
360 }
362 return fix_up(tree);
363 }
365 /*static struct rbnode *find(struct rbtree *rb, struct rbnode *node, void *key)
366 {
367 int cmp;
369 if(!node)
370 return 0;
372 if((cmp = rb->cmp(key, node->key)) == 0) {
373 return node;
374 }
375 return find(rb, cmp < 0 ? node->left : node->right, key);
376 }*/
378 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls)
379 {
380 if(!node)
381 return;
383 traverse(node->left, func, cls);
384 func(node, cls);
385 traverse(node->right, func, cls);
386 }
388 /* helpers */
390 static int is_red(struct rbnode *tree)
391 {
392 return tree && tree->red;
393 }
395 static void color_flip(struct rbnode *tree)
396 {
397 tree->red = !tree->red;
398 tree->left->red = !tree->left->red;
399 tree->right->red = !tree->right->red;
400 }
402 static struct rbnode *rot_left(struct rbnode *a)
403 {
404 struct rbnode *b = a->right;
405 a->right = b->left;
406 b->left = a;
407 b->red = a->red;
408 a->red = 1;
409 return b;
410 }
412 static struct rbnode *rot_right(struct rbnode *a)
413 {
414 struct rbnode *b = a->left;
415 a->left = b->right;
416 b->right = a;
417 b->red = a->red;
418 a->red = 1;
419 return b;
420 }
422 static struct rbnode *find_min(struct rbnode *tree)
423 {
424 if(!tree)
425 return 0;
427 while(tree->left) {
428 tree = tree->left;
429 }
430 return tree;
431 }
433 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree)
434 {
435 if(!tree->left) {
436 if(rb->del) {
437 rb->del(tree->left, rb->del_cls);
438 }
439 rb->free(tree->left);
440 return 0;
441 }
443 /* make sure we've got red (3/4-nodes) at the left side so we can delete at the bottom */
444 if(!is_red(tree->left) && !is_red(tree->left->left)) {
445 tree = move_red_left(tree);
446 }
447 tree->left = del_min(rb, tree->left);
449 /* fix right-reds, red-reds, and split 4-nodes on the way up */
450 return fix_up(tree);
451 }
453 #if 0
454 /* push a red link on this node to the right */
455 static struct rbnode *move_red_right(struct rbnode *tree)
456 {
457 /* flipping it makes both children go red, so we have a red to the right */
458 color_flip(tree);
460 /* if after the flip we've got a red-red situation to the left, fix it */
461 if(is_red(tree->left->left)) {
462 tree = rot_right(tree);
463 color_flip(tree);
464 }
465 return tree;
466 }
467 #endif
469 /* push a red link on this node to the left */
470 static struct rbnode *move_red_left(struct rbnode *tree)
471 {
472 /* flipping it makes both children go red, so we have a red to the left */
473 color_flip(tree);
475 /* if after the flip we've got a red-red on the right-left, fix it */
476 if(is_red(tree->right->left)) {
477 tree->right = rot_right(tree->right);
478 tree = rot_left(tree);
479 color_flip(tree);
480 }
481 return tree;
482 }
484 static struct rbnode *fix_up(struct rbnode *tree)
485 {
486 /* fix right-leaning */
487 if(is_red(tree->right)) {
488 tree = rot_left(tree);
489 }
490 /* change invalid red-red pairs into a proper 4-node */
491 if(is_red(tree->left) && is_red(tree->left->left)) {
492 tree = rot_right(tree);
493 }
494 /* split 4-nodes */
495 if(is_red(tree->left) && is_red(tree->right)) {
496 color_flip(tree);
497 }
498 return tree;
499 }