2 This file is part of ``kdtree'', a library for working with kd-trees.
3 Copyright (C) 2007-2011 John Tsiombikas <nuclear@member.fsf.org>
5 Redistribution and use in source and binary forms, with or without
6 modification, are permitted provided that the following conditions are met:
8 1. Redistributions of source code must retain the above copyright notice, this
9 list of conditions and the following disclaimer.
10 2. Redistributions in binary form must reproduce the above copyright notice,
11 this list of conditions and the following disclaimer in the documentation
12 and/or other materials provided with the distribution.
13 3. The name of the author may not be used to endorse or promote products
14 derived from this software without specific prior written permission.
16 THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR IMPLIED
17 WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
18 MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
19 EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20 EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
21 OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
22 INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
23 CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
24 IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
27 /* single nearest neighbor search written by Tamas Nepusz <tamas@cs.rhul.ac.uk> */
34 #if defined(WIN32) || defined(__WIN32__)
38 #ifdef USE_LIST_NODE_ALLOCATOR
44 #ifndef I_WANT_THREAD_BUGS
45 #error "You are compiling with the fast list node allocator, with pthreads disabled! This WILL break if used from multiple threads."
46 #endif /* I want thread bugs */
48 #endif /* pthread support */
49 #endif /* use list node allocator */
53 double *min, *max; /* minimum/maximum coords */
61 struct kdnode *left, *right; /* negative/positive side */
67 struct res_node *next;
73 struct kdhyperrect *rect;
79 struct res_node *rlist, *riter;
83 #define SQ(x) ((x) * (x))
86 static void clear_rec(struct kdnode *node, void (*destr)(void*));
87 static int insert_rec(struct kdnode **node, const double *pos, void *data, int dir, int dim);
88 static int rlist_insert(struct res_node *list, struct kdnode *item, double dist_sq);
89 static void clear_results(struct kdres *set);
91 static struct kdhyperrect* hyperrect_create(int dim, const double *min, const double *max);
92 static void hyperrect_free(struct kdhyperrect *rect);
93 static struct kdhyperrect* hyperrect_duplicate(const struct kdhyperrect *rect);
94 static void hyperrect_extend(struct kdhyperrect *rect, const double *pos);
95 static double hyperrect_dist_sq(struct kdhyperrect *rect, const double *pos);
97 #ifdef USE_LIST_NODE_ALLOCATOR
98 static struct res_node *alloc_resnode(void);
99 static void free_resnode(struct res_node*);
101 #define alloc_resnode() malloc(sizeof(struct res_node))
102 #define free_resnode(n) free(n)
107 struct kdtree *kd_create(int k)
111 if(!(tree = malloc(sizeof *tree))) {
123 void kd_free(struct kdtree *tree)
131 static void clear_rec(struct kdnode *node, void (*destr)(void*))
135 clear_rec(node->left, destr);
136 clear_rec(node->right, destr);
145 void kd_clear(struct kdtree *tree)
147 clear_rec(tree->root, tree->destr);
151 hyperrect_free(tree->rect);
156 void kd_data_destructor(struct kdtree *tree, void (*destr)(void*))
162 static int insert_rec(struct kdnode **nptr, const double *pos, void *data, int dir, int dim)
168 if(!(node = malloc(sizeof *node))) {
171 if(!(node->pos = malloc(dim * sizeof *node->pos))) {
175 memcpy(node->pos, pos, dim * sizeof *node->pos);
178 node->left = node->right = 0;
184 new_dir = (node->dir + 1) % dim;
185 if(pos[node->dir] < node->pos[node->dir]) {
186 return insert_rec(&(*nptr)->left, pos, data, new_dir, dim);
188 return insert_rec(&(*nptr)->right, pos, data, new_dir, dim);
191 int kd_insert(struct kdtree *tree, const double *pos, void *data)
193 if (insert_rec(&tree->root, pos, data, 0, tree->dim)) {
197 if (tree->rect == 0) {
198 tree->rect = hyperrect_create(tree->dim, pos, pos);
200 hyperrect_extend(tree->rect, pos);
206 int kd_insertf(struct kdtree *tree, const float *pos, void *data)
208 static double sbuf[16];
209 double *bptr, *buf = 0;
210 int res, dim = tree->dim;
215 bptr = buf = alloca(dim * sizeof *bptr);
218 if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
229 res = kd_insert(tree, buf, data);
239 int kd_insert3(struct kdtree *tree, double x, double y, double z, void *data)
245 return kd_insert(tree, buf, data);
248 int kd_insert3f(struct kdtree *tree, float x, float y, float z, void *data)
254 return kd_insert(tree, buf, data);
257 static int find_nearest(struct kdnode *node, const double *pos, double range, struct res_node *list, int ordered, int dim)
260 int i, ret, added_res = 0;
265 for(i=0; i<dim; i++) {
266 dist_sq += SQ(node->pos[i] - pos[i]);
268 if(dist_sq <= SQ(range)) {
269 if(rlist_insert(list, node, ordered ? dist_sq : -1.0) == -1) {
275 dx = pos[node->dir] - node->pos[node->dir];
277 ret = find_nearest(dx <= 0.0 ? node->left : node->right, pos, range, list, ordered, dim);
278 if(ret >= 0 && fabs(dx) < range) {
280 ret = find_nearest(dx <= 0.0 ? node->right : node->left, pos, range, list, ordered, dim);
291 static int find_nearest_n(struct kdnode *node, const double *pos, double range, int num, struct rheap *heap, int dim)
294 int i, ret, added_res = 0;
298 /* if the photon is close enough, add it to the result heap */
300 for(i=0; i<dim; i++) {
301 dist_sq += SQ(node->pos[i] - pos[i]);
303 if(dist_sq <= range_sq) {
304 if(heap->size >= num) {
305 /* get furthest element */
306 struct res_node *maxelem = rheap_get_max(heap);
308 /* and check if the new one is closer than that */
309 if(maxelem->dist_sq > dist_sq) {
310 rheap_remove_max(heap);
312 if(rheap_insert(heap, node, dist_sq) == -1) {
320 if(rheap_insert(heap, node, dist_sq) == -1) {
328 /* find signed distance from the splitting plane */
329 dx = pos[node->dir] - node->pos[node->dir];
331 ret = find_nearest_n(dx <= 0.0 ? node->left : node->right, pos, range, num, heap, dim);
332 if(ret >= 0 && fabs(dx) < range) {
334 ret = find_nearest_n(dx <= 0.0 ? node->right : node->left, pos, range, num, heap, dim);
340 static void kd_nearest_i(struct kdnode *node, const double *pos, struct kdnode **result, double *result_dist_sq, struct kdhyperrect* rect)
344 double dummy, dist_sq;
345 struct kdnode *nearer_subtree, *farther_subtree;
346 double *nearer_hyperrect_coord, *farther_hyperrect_coord;
348 /* Decide whether to go left or right in the tree */
349 dummy = pos[dir] - node->pos[dir];
351 nearer_subtree = node->left;
352 farther_subtree = node->right;
353 nearer_hyperrect_coord = rect->max + dir;
354 farther_hyperrect_coord = rect->min + dir;
356 nearer_subtree = node->right;
357 farther_subtree = node->left;
358 nearer_hyperrect_coord = rect->min + dir;
359 farther_hyperrect_coord = rect->max + dir;
362 if (nearer_subtree) {
363 /* Slice the hyperrect to get the hyperrect of the nearer subtree */
364 dummy = *nearer_hyperrect_coord;
365 *nearer_hyperrect_coord = node->pos[dir];
366 /* Recurse down into nearer subtree */
367 kd_nearest_i(nearer_subtree, pos, result, result_dist_sq, rect);
369 *nearer_hyperrect_coord = dummy;
372 /* Check the distance of the point at the current node, compare it
373 * with our best so far */
375 for(i=0; i < rect->dim; i++) {
376 dist_sq += SQ(node->pos[i] - pos[i]);
378 if (dist_sq < *result_dist_sq) {
380 *result_dist_sq = dist_sq;
383 if (farther_subtree) {
384 /* Get the hyperrect of the farther subtree */
385 dummy = *farther_hyperrect_coord;
386 *farther_hyperrect_coord = node->pos[dir];
387 /* Check if we have to recurse down by calculating the closest
388 * point of the hyperrect and see if it's closer than our
389 * minimum distance in result_dist_sq. */
390 if (hyperrect_dist_sq(rect, pos) < *result_dist_sq) {
391 /* Recurse down into farther subtree */
392 kd_nearest_i(farther_subtree, pos, result, result_dist_sq, rect);
394 /* Undo the slice on the hyperrect */
395 *farther_hyperrect_coord = dummy;
399 struct kdres *kd_nearest(struct kdtree *kd, const double *pos)
401 struct kdhyperrect *rect;
402 struct kdnode *result;
408 if (!kd->rect) return 0;
410 /* Allocate result set */
411 if(!(rset = malloc(sizeof *rset))) {
414 if(!(rset->rlist = alloc_resnode())) {
418 rset->rlist->next = 0;
421 /* Duplicate the bounding hyperrectangle, we will work on the copy */
422 if (!(rect = hyperrect_duplicate(kd->rect))) {
427 /* Our first guesstimate is the root node */
430 for (i = 0; i < kd->dim; i++)
431 dist_sq += SQ(result->pos[i] - pos[i]);
433 /* Search for the nearest neighbour recursively */
434 kd_nearest_i(kd->root, pos, &result, &dist_sq, rect);
436 /* Free the copy of the hyperrect */
437 hyperrect_free(rect);
439 /* Store the result */
441 if (rlist_insert(rset->rlist, result, -1.0) == -1) {
454 struct kdres *kd_nearestf(struct kdtree *tree, const float *pos)
456 static double sbuf[16];
457 double *bptr, *buf = 0;
464 bptr = buf = alloca(dim * sizeof *bptr);
467 if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
478 res = kd_nearest(tree, buf);
488 struct kdres *kd_nearest3(struct kdtree *tree, double x, double y, double z)
494 return kd_nearest(tree, pos);
497 struct kdres *kd_nearest3f(struct kdtree *tree, float x, float y, float z)
503 return kd_nearest(tree, pos);
506 /* ---- nearest N search ---- */
508 static kdres *kd_nearest_n(struct kdtree *kd, const double *pos, int num)
513 if(!(rset = malloc(sizeof *rset))) {
516 if(!(rset->rlist = alloc_resnode())) {
520 rset->rlist->next = 0;
523 if((ret = find_nearest_n(kd->root, pos, range, num, rset->rlist, kd->dim)) == -1) {
532 struct kdres *kd_nearest_range(struct kdtree *kd, const double *pos, double range)
537 if(!(rset = malloc(sizeof *rset))) {
540 if(!(rset->rlist = alloc_resnode())) {
544 rset->rlist->next = 0;
547 if((ret = find_nearest(kd->root, pos, range, rset->rlist, 0, kd->dim)) == -1) {
556 struct kdres *kd_nearest_rangef(struct kdtree *kd, const float *pos, float range)
558 static double sbuf[16];
559 double *bptr, *buf = 0;
566 bptr = buf = alloca(dim * sizeof *bptr);
569 if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
580 res = kd_nearest_range(kd, buf, range);
590 struct kdres *kd_nearest_range3(struct kdtree *tree, double x, double y, double z, double range)
596 return kd_nearest_range(tree, buf, range);
599 struct kdres *kd_nearest_range3f(struct kdtree *tree, float x, float y, float z, float range)
605 return kd_nearest_range(tree, buf, range);
608 void kd_res_free(struct kdres *rset)
611 free_resnode(rset->rlist);
615 int kd_res_size(struct kdres *set)
620 void kd_res_rewind(struct kdres *rset)
622 rset->riter = rset->rlist->next;
625 int kd_res_end(struct kdres *rset)
627 return rset->riter == 0;
630 int kd_res_next(struct kdres *rset)
632 rset->riter = rset->riter->next;
633 return rset->riter != 0;
636 void *kd_res_item(struct kdres *rset, double *pos)
640 memcpy(pos, rset->riter->item->pos, rset->tree->dim * sizeof *pos);
642 return rset->riter->item->data;
647 void *kd_res_itemf(struct kdres *rset, float *pos)
652 for(i=0; i<rset->tree->dim; i++) {
653 pos[i] = rset->riter->item->pos[i];
656 return rset->riter->item->data;
661 void *kd_res_item3(struct kdres *rset, double *x, double *y, double *z)
664 if(*x) *x = rset->riter->item->pos[0];
665 if(*y) *y = rset->riter->item->pos[1];
666 if(*z) *z = rset->riter->item->pos[2];
671 void *kd_res_item3f(struct kdres *rset, float *x, float *y, float *z)
674 if(*x) *x = rset->riter->item->pos[0];
675 if(*y) *y = rset->riter->item->pos[1];
676 if(*z) *z = rset->riter->item->pos[2];
681 void *kd_res_item_data(struct kdres *set)
683 return kd_res_item(set, 0);
686 /* ---- hyperrectangle helpers ---- */
687 static struct kdhyperrect* hyperrect_create(int dim, const double *min, const double *max)
689 size_t size = dim * sizeof(double);
690 struct kdhyperrect* rect = 0;
692 if (!(rect = malloc(sizeof(struct kdhyperrect)))) {
697 if (!(rect->min = malloc(size))) {
701 if (!(rect->max = malloc(size))) {
706 memcpy(rect->min, min, size);
707 memcpy(rect->max, max, size);
712 static void hyperrect_free(struct kdhyperrect *rect)
719 static struct kdhyperrect* hyperrect_duplicate(const struct kdhyperrect *rect)
721 return hyperrect_create(rect->dim, rect->min, rect->max);
724 static void hyperrect_extend(struct kdhyperrect *rect, const double *pos)
728 for (i=0; i < rect->dim; i++) {
729 if (pos[i] < rect->min[i]) {
730 rect->min[i] = pos[i];
732 if (pos[i] > rect->max[i]) {
733 rect->max[i] = pos[i];
738 static double hyperrect_dist_sq(struct kdhyperrect *rect, const double *pos)
743 for (i=0; i < rect->dim; i++) {
744 if (pos[i] < rect->min[i]) {
745 result += SQ(rect->min[i] - pos[i]);
746 } else if (pos[i] > rect->max[i]) {
747 result += SQ(rect->max[i] - pos[i]);
754 /* ---- static helpers ---- */
756 #ifdef USE_LIST_NODE_ALLOCATOR
757 /* special list node allocators. */
758 static struct res_node *free_nodes;
761 static pthread_mutex_t alloc_mutex = PTHREAD_MUTEX_INITIALIZER;
764 static struct res_node *alloc_resnode(void)
766 struct res_node *node;
769 pthread_mutex_lock(&alloc_mutex);
773 node = malloc(sizeof *node);
776 free_nodes = free_nodes->next;
781 pthread_mutex_unlock(&alloc_mutex);
787 static void free_resnode(struct res_node *node)
790 pthread_mutex_lock(&alloc_mutex);
793 node->next = free_nodes;
797 pthread_mutex_unlock(&alloc_mutex);
800 #endif /* list node allocator or not */
803 /* inserts the item. if dist_sq is >= 0, then do an ordered insert */
804 /* TODO make the ordering code use heapsort */
805 static int rlist_insert(struct res_node *list, struct kdnode *item, double dist_sq)
807 struct res_node *rnode;
809 if(!(rnode = alloc_resnode())) {
813 rnode->dist_sq = dist_sq;
816 while(list->next && list->next->dist_sq < dist_sq) {
820 rnode->next = list->next;
825 static void clear_results(struct kdres *rset)
827 struct res_node *tmp, *node = rset->rlist->next;
835 rset->rlist->next = 0;