diff options
Diffstat (limited to 'src/common/libalgo/btree.c')
-rw-r--r-- | src/common/libalgo/btree.c | 295 |
1 files changed, 295 insertions, 0 deletions
diff --git a/src/common/libalgo/btree.c b/src/common/libalgo/btree.c new file mode 100644 index 0000000..da0e25b --- /dev/null +++ b/src/common/libalgo/btree.c @@ -0,0 +1,295 @@ +#include <malloc.h> +#include <debug.h> + +#include <btree.h> + +typedef struct btree_item { + void *key, *val; + + int height; + struct btree_item *left, *right; +} btree_item_t; + +struct btree { + key_cmp_fun_t cf; + kv_iter_fun_t releasef; + + btree_item_t *root; + int nitems; +}; + +btree_t* create_btree(key_cmp_fun_t cf, kv_iter_fun_t relf) { + btree_t* t = (btree_t*)malloc(sizeof(btree_t)); + + if (t == 0) return 0; + + t->cf = cf; + t->releasef = relf; + + t->root = 0; + t->nitems = 0; + + return t; +} + +void delete_btree(btree_t *t) { + void delete_item_rec(btree_item_t *i,kv_iter_fun_t relf) { + if (i == 0) return; + + delete_item_rec(i->left, relf); + delete_item_rec(i->right, relf); + + relf(i->key, i->val); + free(i); + } + + delete_item_rec(t->root, t->releasef); + + free(t); +} + +size_t btree_count(btree_t *i) { + return i->nitems; +} + +// ============== // +// TREE ROTATIONS // +// ============== // + +static int height(btree_item_t *i) { + if (i == 0) return 0; + return i->height; +} + +static void recalc_height(btree_item_t *i) { + if (i == 0) return; + + i->height = MAX(height(i->left), height(i->right)) + 1; +} + +static btree_item_t* rotate_left(btree_item_t *i) { + // I(X(a, b), Y) -> a X b I Y -> X(a, I(b, Y)) + + btree_item_t *x = i->left; + if (x == 0) return i; + btree_item_t *b = x->right; + + x->right = i; + i->left = b; + + recalc_height(i); + recalc_height(x); + + return x; +} + +static btree_item_t* rotate_right(btree_item_t *i) { + // I(X, Y(a,b)) -> X I a Y b -> Y(I(X, a), b) + + btree_item_t *y = i->right; + if (y == 0) return i; + btree_item_t *a = y->left; + + y->left = i; + i->right = a; + + recalc_height(i); + recalc_height(y); + + return y; +} + +static btree_item_t* equilibrate(btree_item_t *i) { + if (i == 0) return 0; + + while (height(i->left) - height(i->right) >= 2) + i = rotate_left(i); + + while (height(i->right) - height(i->left) >= 2) + i = rotate_right(i); + + return i; +} + +// =================== // +// ADDING AND DELETING // +// =================== // + +bool btree_add(btree_t *t, void* key, void* val) { + btree_item_t* insert(btree_t *t, btree_item_t *root, btree_item_t *i) { + if (root == 0) return i; + + if (t->cf(i->key, root->key) <= 0) { + root->left = insert(t, root->left, i); + } else { + root->right = insert(t, root->right, i); + } + recalc_height(root); + + return equilibrate(root); + } + + btree_item_t *x = (btree_item_t*)malloc(sizeof(btree_item_t)); + if (x == 0) return false; + + x->key = key; + x->val = val; + x->left = x->right = 0; + recalc_height(x); + + t->root = insert(t, t->root, x); + t->nitems++; + + return true; +} + +void btree_remove(btree_t *t, const void* key) { + btree_item_t *extract_smallest(btree_item_t *i, btree_item_t **out_smallest) { + ASSERT(i != 0); + if (i->left == 0) { + *out_smallest = i; + return i->right; + } else { + i->left = extract_smallest(i->left, out_smallest); + recalc_height(i); + return equilibrate(i); + } + } + + btree_item_t *remove_aux(btree_t *t, btree_item_t *i, const void* key) { + if (i == 0) return 0; + + int c = t->cf(key, i->key); + if (c < 0) { + i->left = remove_aux(t, i->left, key); + return equilibrate(i); + } else if (c > 0) { + i->right = remove_aux(t, i->right, key); + return equilibrate(i); + } else { + // remove item i + + btree_item_t *new_i; + if (i->right == 0 || i->left == 0) { + new_i = (i->right == 0 ? i->left : i->right); + } else { + btree_item_t *new_i_right = extract_smallest(i->right, &new_i); + + new_i->left = i->left; + new_i->right = new_i_right; + + recalc_height(new_i); + new_i = equilibrate(new_i); + } + + t->releasef(i->key, i->val); + free(i); + t->nitems--; + + return remove_aux(t, new_i, key); // loop because several elements may correspond + } + } + + t->root = remove_aux(t, t->root, key); +} + +// ======================== // +// LOOKING UP AND ITERATING // +// ======================== // + +void* btree_find(btree_t *t, const void* key) { + btree_item_t *find_aux(btree_t *t, btree_item_t *i, const void* key) { + if (i == 0) return 0; + + int c = t->cf(key, i->key); + if (c == 0) { + return i; + } else if (c < 0) { + return find_aux(t, i->left, key); + } else { + return find_aux(t, i->right, key); + } + } + + btree_item_t *i = find_aux(t, t->root, key); + + if (i == 0) return 0; + return i->val; +} + +void* btree_lower(btree_t *t, const void* key) { + btree_item_t *find_aux(btree_t *t, btree_item_t *i, const void* key) { + if (i == 0) return 0; + + int c = t->cf(key, i->key); + if (c == 0) { + return i; + } else if (c < 0) { + return find_aux(t, i->left, key); + } else { + btree_item_t *r = find_aux(t, i->right, key); + if (r == 0) r = i; + return r; + } + } + + btree_item_t *i = find_aux(t, t->root, key); + + if (i == 0) return 0; + return i->val; +} + +void* btree_upper(btree_t *t, const void* key) { + btree_item_t *find_aux(btree_t *t, btree_item_t *i, const void* key) { + if (i == 0) return 0; + + int c = t->cf(key, i->key); + if (c == 0) { + return i; + } else if (c < 0) { + btree_item_t *r = find_aux(t, i->left, key); + if (r == 0) r = i; + return r; + } else { + return find_aux(t, i->right, key); + } + } + + btree_item_t *i = find_aux(t, t->root, key); + + if (i == 0) return 0; + return i->val; +} + +void btree_iter(btree_t *t, kv_iter_fun_t f) { + void iter_aux(btree_item_t *i, kv_iter_fun_t f) { + if (i == 0) return; + + iter_aux(i->left, f); + f(i->key, i->val); + iter_aux(i->right, f); + } + + iter_aux(t->root, f); +} + +void btree_iter_on(btree_t *t, const void* key, kv_iter_fun_t f) { + void iter_aux(btree_t *t, btree_item_t *i, const void* key, kv_iter_fun_t f) { + if (i == 0) return; + + int c = t->cf(key, i->key); + if (c == 0) { + iter_aux(t, i->left, key, f); + f(i->key, i->val); + iter_aux(t, i->right, key, f); + } else if (c < 0) { + iter_aux(t, i->left, key, f); + } else { + iter_aux(t, i->right, key, f); + } + } + + iter_aux(t, t->root, key, f); +} + +/* vim: set ts=4 sw=4 tw=0 noet :*/ |