diff --git a/examples/Makefile b/examples/Makefile index dab0f7e..ea683e4 100644 --- a/examples/Makefile +++ b/examples/Makefile @@ -5,10 +5,11 @@ CFLAGS = -std=c89 -pedantic -Wall -g -I.. LDFLAGS = $(kdlib) -lm .PHONY: all -all: test test2 +all: test test2 test3 test: test.c $(kdlib) test2: test2.c $(kdlib) +test3: test3.c $(kdlib) .PHONY: clean clean: diff --git a/examples/test.c b/examples/test.c index 89e6195..76defd2 100644 --- a/examples/test.c +++ b/examples/test.c @@ -23,7 +23,7 @@ unsigned int get_msec(void) int main(int argc, char **argv) { - int i, vcount = 10; + int i, vcount = 10000; void *kd, *set; unsigned int msec, start; @@ -52,7 +52,7 @@ int main(int argc, char **argv) msec = get_msec() - start; printf("range query returned %d items in %.5f sec\n", kd_res_size(set), (float)msec / 1000.0); kd_res_free(set); - + kd_clear(kd); kd_free(kd); return 0; } diff --git a/examples/test2.c b/examples/test2.c index f06859a..ae6c3d1 100644 --- a/examples/test2.c +++ b/examples/test2.c @@ -8,7 +8,7 @@ #include #include "kdtree.h" -#define DEF_NUM_PTS 10 +#define DEF_NUM_PTS 10000 /* returns the distance squared between two dims-dimensional double arrays */ static double dist_sq( double *a1, double *a2, int dims ); diff --git a/examples/test3.c b/examples/test3.c new file mode 100644 index 0000000..88a09fc --- /dev/null +++ b/examples/test3.c @@ -0,0 +1,94 @@ +/*! gcc -std=c89 -pedantic -Wall -g -o test2 test2.c libkdtree.a -lm */ +/* Extended test program, contributed by David Underhill */ +#include +#include +#include +#include +#include +#include +#include "kdtree.h" + +#define DEF_NUM_PTS 10000 +#define DIM 3 + +/* returns the distance squared between two dims-dimensional double arrays */ +static double dist_sq( double *a1, double *a2, int dims ); + +/* get a random double between -10 and 10 */ +static double rd( void ); + +int main(int argc, char **argv) { + int i, num_pts = DEF_NUM_PTS; + void *ptree; + char *data, *pch; + struct kdres *presults; + double pos[3], dist, minpos[3], maxpos[3]; + double pt[3] = { 0, 0, 1 }; + double radius = 10; + + if(argc > 1 && isdigit(argv[1][0])) { + num_pts = atoi(argv[1]); + } + + if(!(data = malloc(num_pts))) { + perror("malloc failed"); + return 1; + } + + srand( time(0) ); + + /* create a k-d tree for 3-dimensional points */ + ptree = kd_create( DIM ); + + /* add some random nodes to the tree (assert nodes are successfully inserted) */ + for( i=0; i= 0 ) { + diff = (a1[dims] - a2[dims]); + dist_sq += diff*diff; + } + return dist_sq; +} + +static double rd( void ) { + return (double)rand()/RAND_MAX * 20.0 - 10.0; +} diff --git a/kdtree.c b/kdtree.c index 41cd9c0..4bffe08 100644 --- a/kdtree.c +++ b/kdtree.c @@ -25,6 +25,7 @@ IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */ /* single nearest neighbor search written by Tamas Nepusz */ +/* in_bounds written by Philip Kovac */ #ifdef HAVE_CONFIG_H #include @@ -63,7 +64,7 @@ struct kdnode { int dir; void *data; - struct kdnode *left, *right; /* negative/positive side */ + struct kdnode *left, *right, *next; /* negative/positive/equal side */ }; struct res_node { @@ -76,6 +77,10 @@ struct kdtree { int dim; struct kdnode *root; struct kdhyperrect *rect; + + struct kdnode** node_stack; + size_t depth, node_stack_max; + int node_stack_head; void (*destr)(void*); }; @@ -83,13 +88,14 @@ struct kdres { struct kdtree *tree; struct res_node *rlist, *riter; int size; + int expanded; }; #define SQ(x) ((x) * (x)) -static void clear_rec(struct kdnode *node, void (*destr)(void*)); -static int insert_rec(struct kdnode **node, const double *pos, void *data, int dir, int dim); +static void clear_rec(struct kdtree *tree, void (*destr)(void*)); +static int insert_rec(struct kdnode **node, const double *pos, void *data, int dim); static int rlist_insert(struct res_node *list, struct kdnode *item, double dist_sq); static void clear_results(struct kdres *set); @@ -107,7 +113,29 @@ static void free_resnode(struct res_node*); #define free_resnode(n) free(n) #endif +static inline void node_stack_prep(struct kdtree* tree) +{ + if (!tree->node_stack || tree->node_stack_max < tree->depth) + { + if (tree->node_stack) + free(tree->node_stack); + tree->node_stack = malloc(sizeof(struct kdnode*) * tree->depth); + tree->node_stack_max = tree->depth; + tree->node_stack_head = 0; + } +} +static inline struct kdnode* node_stack_pop(struct kdtree* tree) +{ + if (tree->node_stack_head == 0) + return NULL; + else + return tree->node_stack[--tree->node_stack_head]; +} +static inline void node_stack_push(struct kdtree* tree, struct kdnode* node) +{ + tree->node_stack[tree->node_stack_head++] = node; +} struct kdtree *kd_create(int k) { @@ -121,6 +149,10 @@ struct kdtree *kd_create(int k) tree->root = 0; tree->destr = 0; tree->rect = 0; + tree->depth = 0; + tree->node_stack = 0; + tree->node_stack_max = 0; + tree->node_stack_head = 0; return tree; } @@ -133,24 +165,38 @@ void kd_free(struct kdtree *tree) } } -static void clear_rec(struct kdnode *node, void (*destr)(void*)) +static void clear_rec(struct kdtree *tree, void (*destr)(void*)) { - if(!node) return; + struct kdnode* node; - clear_rec(node->left, destr); - clear_rec(node->right, destr); - - if(destr) { - destr(node->data); + node_stack_prep(tree); + + while ((node = node_stack_pop(tree))) + { + if (node->left) node_stack_push(tree, node->left); + if (node->next) node_stack_push(tree, node->next); + if (node->right) node_stack_push(tree, node->right); + + if (destr) { + destr(node->data); + }; + free(node->pos); + free(node); } - free(node->pos); - free(node); } void kd_clear(struct kdtree *tree) { - clear_rec(tree->root, tree->destr); + clear_rec(tree, tree->destr); tree->root = 0; + tree->depth = 0; + + if (tree->node_stack) + { + free(tree->node_stack); + tree->node_stack = 0; + tree->node_stack_head = tree->node_stack_max = 0; + } if (tree->rect) { hyperrect_free(tree->rect); @@ -164,40 +210,74 @@ void kd_data_destructor(struct kdtree *tree, void (*destr)(void*)) } -static int insert_rec(struct kdnode **nptr, const double *pos, void *data, int dir, int dim) +static int insert_rec(struct kdnode **nptr, const double *pos, void *data, int dim) { - int new_dir; struct kdnode *node; - - if(!*nptr) { - if(!(node = malloc(sizeof *node))) { - return -1; + size_t depth = 0; + int pos_idx, dir, new_dir = 0; + while (*nptr) + { + depth++; + node = *nptr; + dir = node->dir; + new_dir = (dir + 1) % dim; + + if (pos[dir] < node->pos[dir]) + { + nptr = &(node->left); + continue; } - if(!(node->pos = malloc(dim * sizeof *node->pos))) { - free(node); - return -1; + else if (pos[dir] == node->pos[dir]) + { + for (pos_idx = 0; pos_idx < dim; pos_idx++) + if (pos[pos_idx] != node->pos[pos_idx]) break; + if (pos_idx == dim) + { + while (node->next) + node = node->next; + nptr = &(node->next); + break; + } } - memcpy(node->pos, pos, dim * sizeof *node->pos); - node->data = data; - node->dir = dir; - node->left = node->right = 0; - *nptr = node; - return 0; + else /* pos[dir] > node->pos[dir] */ + { + nptr = &(node->right); + continue; + } + } + + if(!(node = malloc(sizeof *node))) { + return -1; } - node = *nptr; - new_dir = (node->dir + 1) % dim; - if(pos[node->dir] < node->pos[node->dir]) { - return insert_rec(&(*nptr)->left, pos, data, new_dir, dim); + if(!(node->pos = malloc(dim * sizeof *node->pos))) { + free(node); + return -1; } - return insert_rec(&(*nptr)->right, pos, data, new_dir, dim); + + memcpy(node->pos, pos, dim * sizeof *node->pos); + + node->data = data; + node->dir = new_dir; + node->left = node->right = node->next = 0; + *nptr = node; + + return depth + 1; } int kd_insert(struct kdtree *tree, const double *pos, void *data) { - if (insert_rec(&tree->root, pos, data, 0, tree->dim)) { + int result; + + result = insert_rec(&tree->root, pos, data, tree->dim); + + if (result == -1) { return -1; } + else + { + if (result > tree->depth) tree->depth = result; + } if (tree->rect == 0) { tree->rect = hyperrect_create(tree->dim, pos, pos); @@ -259,6 +339,55 @@ int kd_insert3f(struct kdtree *tree, float x, float y, float z, void *data) return kd_insert(tree, buf, data); } +static int in_bounds(struct kdtree* tree, const double* min_pos, const double* max_pos, struct res_node *list, int ordered, int dim, int inclusive) +{ + int i, is_in_bounds, added_res = 0; + struct kdnode* node; + + node_stack_prep(tree); + node_stack_push(tree, tree->root); + + while ((node = node_stack_pop(tree))) + { + is_in_bounds = 1; + + if (inclusive) + { + for (i=0; ipos[i] < min_pos[i] || node->pos[i] > max_pos[i]) + { + is_in_bounds = 0; + break; + } + } + else + { + for (i=0; ipos[i] <= min_pos[i] || node->pos[i] >= max_pos[i]) + { + is_in_bounds = 0; + break; + } + } + + if (is_in_bounds) + { + if (rlist_insert(list, node, -1.0) == -1) { + return -1; + } + added_res++; + } + + if (!(min_pos[node->dir] > node->pos[node->dir]) && node->left) + node_stack_push(tree, node->left); + + if (!(max_pos[node->dir] < node->pos[node->dir]) && node->right) + node_stack_push(tree, node->right); + } + + return added_res; +} + static int find_nearest(struct kdnode *node, const double *pos, double range, struct res_node *list, int ordered, int dim) { double dist_sq, dx; @@ -421,6 +550,7 @@ struct kdres *kd_nearest(struct kdtree *kd, const double *pos) return 0; } rset->rlist->next = 0; + rset->expanded = 0; rset->tree = kd; /* Duplicate the bounding hyperrectangle, we will work on the copy */ @@ -523,6 +653,7 @@ static kdres *kd_nearest_n(struct kdtree *kd, const double *pos, int num) return 0; } rset->rlist->next = 0; + rset->expanded = 0; rset->tree = kd; if((ret = find_nearest_n(kd->root, pos, range, num, rset->rlist, kd->dim)) == -1) { @@ -534,6 +665,34 @@ static kdres *kd_nearest_n(struct kdtree *kd, const double *pos, int num) return rset; }*/ +struct kdres *kd_in_bounds(struct kdtree *kd, const double* min_pos, const double* max_pos, int inclusive) +{ + int ret; + struct kdres *rset; + + if (!(rset = malloc(sizeof *rset))) { + return 0; + } + + if (!(rset->rlist = alloc_resnode())) { + free(rset); + return 0; + } + + rset->rlist->next = 0; + rset->expanded = 0; + rset->tree = kd; + + if ((ret = in_bounds(kd, min_pos, max_pos, rset->rlist, 0, kd->dim, inclusive)) == -1) { + kd_res_free(rset); + return 0; + } + + rset->size = ret; + kd_res_rewind(rset); + return rset; +} + struct kdres *kd_nearest_range(struct kdtree *kd, const double *pos, double range) { int ret; @@ -547,6 +706,7 @@ struct kdres *kd_nearest_range(struct kdtree *kd, const double *pos, double rang return 0; } rset->rlist->next = 0; + rset->expanded = 0; rset->tree = kd; if((ret = find_nearest(kd->root, pos, range, rset->rlist, 0, kd->dim)) == -1) { @@ -622,9 +782,43 @@ int kd_res_size(struct kdres *set) return (set->size); } +static int res_expand_node(struct res_node** rnode) +{ + int size_inc = 0; + struct kdnode *node; + node = (*rnode)->item; + while (node->next) + { + *rnode = ((*rnode)->next = alloc_resnode()); + + (*rnode)->item = node->next; + node = node->next; + size_inc++; + } + return size_inc; +} + +static int res_expand(struct kdres *rset) +{ + struct res_node *rnode, *next_node; + if (rset->expanded) + return 0; /* Has already been expanded, don't do anything */ + rnode = rset->riter; + while (rnode) + { + next_node = rnode->next; + rset->size += res_expand_node(&rnode); + rnode = (rnode->next = next_node); + } + rset->expanded = 1; + return 0; +} + void kd_res_rewind(struct kdres *rset) { rset->riter = rset->rlist->next; + if (!(rset->expanded)) + res_expand(rset); } int kd_res_end(struct kdres *rset) @@ -838,4 +1032,5 @@ static void clear_results(struct kdres *rset) } rset->rlist->next = 0; + rset->expanded = 0; } diff --git a/kdtree.h b/kdtree.h index 92d43e4..9115269 100644 --- a/kdtree.h +++ b/kdtree.h @@ -93,6 +93,14 @@ struct kdres *kd_nearest_rangef(struct kdtree *tree, const float *pos, float ran struct kdres *kd_nearest_range3(struct kdtree *tree, double x, double y, double z, double range); struct kdres *kd_nearest_range3f(struct kdtree *tree, float x, float y, float z, float range); +/* Find nodes within a given range. + * This function returns a pointer to a result set. + * The returned pointer can be null as an indication of an error. Otherwise + * a valid result set is always returned which may contain 0 or more elements. + * The result set must be deallocated with kd_res_free after use. + */ +struct kdres *kd_in_bounds(struct kdtree *tree, const double *min_pos, const double *max_pos, int inclusive); + /* frees a result set returned by kd_nearest_range() */ void kd_res_free(struct kdres *set);