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> */
39 #if defined(WIN32) || defined(__WIN32__)
43 #ifdef USE_LIST_NODE_ALLOCATOR
49 #ifndef I_WANT_THREAD_BUGS
50 #error "You are compiling with the fast list node allocator, with pthreads disabled! This WILL break if used from multiple threads."
51 #endif /* I want thread bugs */
53 #endif /* pthread support */
54 #endif /* use list node allocator */
58 double *min, *max; /* minimum/maximum coords */
66 struct kdnode *left, *right; /* negative/positive side */
72 struct res_node *next;
78 struct kdhyperrect *rect;
84 struct res_node *rlist, *riter;
88 #define SQ(x) ((x) * (x))
91 static void clear_rec(struct kdnode *node, void (*destr)(void*));
92 static int insert_rec(struct kdnode **node, const double *pos, void *data, int dir, int dim);
93 static int rlist_insert(struct res_node *list, struct kdnode *item, double dist_sq);
94 static void clear_results(struct kdres *set);
96 static struct kdhyperrect* hyperrect_create(int dim, const double *min, const double *max);
97 static void hyperrect_free(struct kdhyperrect *rect);
98 static struct kdhyperrect* hyperrect_duplicate(const struct kdhyperrect *rect);
99 static void hyperrect_extend(struct kdhyperrect *rect, const double *pos);
100 static double hyperrect_dist_sq(struct kdhyperrect *rect, const double *pos);
102 #ifdef USE_LIST_NODE_ALLOCATOR
103 static struct res_node *alloc_resnode(void);
104 static void free_resnode(struct res_node*);
106 #define alloc_resnode() malloc(sizeof(struct res_node))
107 #define free_resnode(n) free(n)
112 struct kdtree *kd_create(int k)
116 if(!(tree = malloc(sizeof *tree))) {
128 void kd_free(struct kdtree *tree)
136 static void clear_rec(struct kdnode *node, void (*destr)(void*))
140 clear_rec(node->left, destr);
141 clear_rec(node->right, destr);
150 void kd_clear(struct kdtree *tree)
152 clear_rec(tree->root, tree->destr);
156 hyperrect_free(tree->rect);
161 void kd_data_destructor(struct kdtree *tree, void (*destr)(void*))
167 static int insert_rec(struct kdnode **nptr, const double *pos, void *data, int dir, int dim)
173 if(!(node = malloc(sizeof *node))) {
176 if(!(node->pos = malloc(dim * sizeof *node->pos))) {
180 memcpy(node->pos, pos, dim * sizeof *node->pos);
183 node->left = node->right = 0;
189 new_dir = (node->dir + 1) % dim;
190 if(pos[node->dir] < node->pos[node->dir]) {
191 return insert_rec(&(*nptr)->left, pos, data, new_dir, dim);
193 return insert_rec(&(*nptr)->right, pos, data, new_dir, dim);
196 int kd_insert(struct kdtree *tree, const double *pos, void *data)
198 if (insert_rec(&tree->root, pos, data, 0, tree->dim)) {
202 if (tree->rect == 0) {
203 tree->rect = hyperrect_create(tree->dim, pos, pos);
205 hyperrect_extend(tree->rect, pos);
211 int kd_insertf(struct kdtree *tree, const float *pos, void *data)
213 static double sbuf[16];
214 double *bptr, *buf = 0;
215 int res, dim = tree->dim;
220 bptr = buf = alloca(dim * sizeof *bptr);
223 if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
234 res = kd_insert(tree, buf, data);
244 int kd_insert3(struct kdtree *tree, double x, double y, double z, void *data)
250 return kd_insert(tree, buf, data);
253 int kd_insert3f(struct kdtree *tree, float x, float y, float z, void *data)
259 return kd_insert(tree, buf, data);
262 static int find_nearest(struct kdnode *node, const double *pos, double range, struct res_node *list, int ordered, int dim)
265 int i, ret, added_res = 0;
270 for(i=0; i<dim; i++) {
271 dist_sq += SQ(node->pos[i] - pos[i]);
273 if(dist_sq <= SQ(range)) {
274 if(rlist_insert(list, node, ordered ? dist_sq : -1.0) == -1) {
280 dx = pos[node->dir] - node->pos[node->dir];
282 ret = find_nearest(dx <= 0.0 ? node->left : node->right, pos, range, list, ordered, dim);
283 if(ret >= 0 && fabs(dx) < range) {
285 ret = find_nearest(dx <= 0.0 ? node->right : node->left, pos, range, list, ordered, dim);
296 static int find_nearest_n(struct kdnode *node, const double *pos, double range, int num, struct rheap *heap, int dim)
299 int i, ret, added_res = 0;
303 /* if the photon is close enough, add it to the result heap */
305 for(i=0; i<dim; i++) {
306 dist_sq += SQ(node->pos[i] - pos[i]);
308 if(dist_sq <= range_sq) {
309 if(heap->size >= num) {
310 /* get furthest element */
311 struct res_node *maxelem = rheap_get_max(heap);
313 /* and check if the new one is closer than that */
314 if(maxelem->dist_sq > dist_sq) {
315 rheap_remove_max(heap);
317 if(rheap_insert(heap, node, dist_sq) == -1) {
325 if(rheap_insert(heap, node, dist_sq) == -1) {
333 /* find signed distance from the splitting plane */
334 dx = pos[node->dir] - node->pos[node->dir];
336 ret = find_nearest_n(dx <= 0.0 ? node->left : node->right, pos, range, num, heap, dim);
337 if(ret >= 0 && fabs(dx) < range) {
339 ret = find_nearest_n(dx <= 0.0 ? node->right : node->left, pos, range, num, heap, dim);
345 static void kd_nearest_i(struct kdnode *node, const double *pos, struct kdnode **result, double *result_dist_sq, struct kdhyperrect* rect)
349 double dummy, dist_sq;
350 struct kdnode *nearer_subtree, *farther_subtree;
351 double *nearer_hyperrect_coord, *farther_hyperrect_coord;
353 /* Decide whether to go left or right in the tree */
354 dummy = pos[dir] - node->pos[dir];
356 nearer_subtree = node->left;
357 farther_subtree = node->right;
358 nearer_hyperrect_coord = rect->max + dir;
359 farther_hyperrect_coord = rect->min + dir;
361 nearer_subtree = node->right;
362 farther_subtree = node->left;
363 nearer_hyperrect_coord = rect->min + dir;
364 farther_hyperrect_coord = rect->max + dir;
367 if (nearer_subtree) {
368 /* Slice the hyperrect to get the hyperrect of the nearer subtree */
369 dummy = *nearer_hyperrect_coord;
370 *nearer_hyperrect_coord = node->pos[dir];
371 /* Recurse down into nearer subtree */
372 kd_nearest_i(nearer_subtree, pos, result, result_dist_sq, rect);
374 *nearer_hyperrect_coord = dummy;
377 /* Check the distance of the point at the current node, compare it
378 * with our best so far */
380 for(i=0; i < rect->dim; i++) {
381 dist_sq += SQ(node->pos[i] - pos[i]);
383 if (dist_sq < *result_dist_sq) {
385 *result_dist_sq = dist_sq;
388 if (farther_subtree) {
389 /* Get the hyperrect of the farther subtree */
390 dummy = *farther_hyperrect_coord;
391 *farther_hyperrect_coord = node->pos[dir];
392 /* Check if we have to recurse down by calculating the closest
393 * point of the hyperrect and see if it's closer than our
394 * minimum distance in result_dist_sq. */
395 if (hyperrect_dist_sq(rect, pos) < *result_dist_sq) {
396 /* Recurse down into farther subtree */
397 kd_nearest_i(farther_subtree, pos, result, result_dist_sq, rect);
399 /* Undo the slice on the hyperrect */
400 *farther_hyperrect_coord = dummy;
404 struct kdres *kd_nearest(struct kdtree *kd, const double *pos)
406 struct kdhyperrect *rect;
407 struct kdnode *result;
413 if (!kd->rect) return 0;
415 /* Allocate result set */
416 if(!(rset = malloc(sizeof *rset))) {
419 if(!(rset->rlist = alloc_resnode())) {
423 rset->rlist->next = 0;
426 /* Duplicate the bounding hyperrectangle, we will work on the copy */
427 if (!(rect = hyperrect_duplicate(kd->rect))) {
432 /* Our first guesstimate is the root node */
435 for (i = 0; i < kd->dim; i++)
436 dist_sq += SQ(result->pos[i] - pos[i]);
438 /* Search for the nearest neighbour recursively */
439 kd_nearest_i(kd->root, pos, &result, &dist_sq, rect);
441 /* Free the copy of the hyperrect */
442 hyperrect_free(rect);
444 /* Store the result */
446 if (rlist_insert(rset->rlist, result, -1.0) == -1) {
459 struct kdres *kd_nearestf(struct kdtree *tree, const float *pos)
461 static double sbuf[16];
462 double *bptr, *buf = 0;
469 bptr = buf = alloca(dim * sizeof *bptr);
472 if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
483 res = kd_nearest(tree, buf);
493 struct kdres *kd_nearest3(struct kdtree *tree, double x, double y, double z)
499 return kd_nearest(tree, pos);
502 struct kdres *kd_nearest3f(struct kdtree *tree, float x, float y, float z)
508 return kd_nearest(tree, pos);
511 /* ---- nearest N search ---- */
513 static kdres *kd_nearest_n(struct kdtree *kd, const double *pos, int num)
518 if(!(rset = malloc(sizeof *rset))) {
521 if(!(rset->rlist = alloc_resnode())) {
525 rset->rlist->next = 0;
528 if((ret = find_nearest_n(kd->root, pos, range, num, rset->rlist, kd->dim)) == -1) {
537 struct kdres *kd_nearest_range(struct kdtree *kd, const double *pos, double range)
542 if(!(rset = malloc(sizeof *rset))) {
545 if(!(rset->rlist = alloc_resnode())) {
549 rset->rlist->next = 0;
552 if((ret = find_nearest(kd->root, pos, range, rset->rlist, 0, kd->dim)) == -1) {
561 struct kdres *kd_nearest_rangef(struct kdtree *kd, const float *pos, float range)
563 static double sbuf[16];
564 double *bptr, *buf = 0;
571 bptr = buf = alloca(dim * sizeof *bptr);
574 if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
585 res = kd_nearest_range(kd, buf, range);
595 struct kdres *kd_nearest_range3(struct kdtree *tree, double x, double y, double z, double range)
601 return kd_nearest_range(tree, buf, range);
604 struct kdres *kd_nearest_range3f(struct kdtree *tree, float x, float y, float z, float range)
610 return kd_nearest_range(tree, buf, range);
613 void kd_res_free(struct kdres *rset)
616 free_resnode(rset->rlist);
620 int kd_res_size(struct kdres *set)
625 void kd_res_rewind(struct kdres *rset)
627 rset->riter = rset->rlist->next;
630 int kd_res_end(struct kdres *rset)
632 return rset->riter == 0;
635 int kd_res_next(struct kdres *rset)
637 rset->riter = rset->riter->next;
638 return rset->riter != 0;
641 void *kd_res_item(struct kdres *rset, double *pos)
645 memcpy(pos, rset->riter->item->pos, rset->tree->dim * sizeof *pos);
647 return rset->riter->item->data;
652 void *kd_res_itemf(struct kdres *rset, float *pos)
657 for(i=0; i<rset->tree->dim; i++) {
658 pos[i] = rset->riter->item->pos[i];
661 return rset->riter->item->data;
666 void *kd_res_item3(struct kdres *rset, double *x, double *y, double *z)
669 if(*x) *x = rset->riter->item->pos[0];
670 if(*y) *y = rset->riter->item->pos[1];
671 if(*z) *z = rset->riter->item->pos[2];
672 return rset->riter->item->data;
677 void *kd_res_item3f(struct kdres *rset, float *x, float *y, float *z)
680 if(*x) *x = rset->riter->item->pos[0];
681 if(*y) *y = rset->riter->item->pos[1];
682 if(*z) *z = rset->riter->item->pos[2];
683 return rset->riter->item->data;
688 void *kd_res_item_data(struct kdres *set)
690 return kd_res_item(set, 0);
693 /* ---- hyperrectangle helpers ---- */
694 static struct kdhyperrect* hyperrect_create(int dim, const double *min, const double *max)
696 size_t size = dim * sizeof(double);
697 struct kdhyperrect* rect = 0;
699 if (!(rect = malloc(sizeof(struct kdhyperrect)))) {
704 if (!(rect->min = malloc(size))) {
708 if (!(rect->max = malloc(size))) {
713 memcpy(rect->min, min, size);
714 memcpy(rect->max, max, size);
719 static void hyperrect_free(struct kdhyperrect *rect)
726 static struct kdhyperrect* hyperrect_duplicate(const struct kdhyperrect *rect)
728 return hyperrect_create(rect->dim, rect->min, rect->max);
731 static void hyperrect_extend(struct kdhyperrect *rect, const double *pos)
735 for (i=0; i < rect->dim; i++) {
736 if (pos[i] < rect->min[i]) {
737 rect->min[i] = pos[i];
739 if (pos[i] > rect->max[i]) {
740 rect->max[i] = pos[i];
745 static double hyperrect_dist_sq(struct kdhyperrect *rect, const double *pos)
750 for (i=0; i < rect->dim; i++) {
751 if (pos[i] < rect->min[i]) {
752 result += SQ(rect->min[i] - pos[i]);
753 } else if (pos[i] > rect->max[i]) {
754 result += SQ(rect->max[i] - pos[i]);
761 /* ---- static helpers ---- */
763 #ifdef USE_LIST_NODE_ALLOCATOR
764 /* special list node allocators. */
765 static struct res_node *free_nodes;
768 static pthread_mutex_t alloc_mutex = PTHREAD_MUTEX_INITIALIZER;
771 static struct res_node *alloc_resnode(void)
773 struct res_node *node;
776 pthread_mutex_lock(&alloc_mutex);
780 node = malloc(sizeof *node);
783 free_nodes = free_nodes->next;
788 pthread_mutex_unlock(&alloc_mutex);
794 static void free_resnode(struct res_node *node)
797 pthread_mutex_lock(&alloc_mutex);
800 node->next = free_nodes;
804 pthread_mutex_unlock(&alloc_mutex);
807 #endif /* list node allocator or not */
810 /* inserts the item. if dist_sq is >= 0, then do an ordered insert */
811 /* TODO make the ordering code use heapsort */
812 static int rlist_insert(struct res_node *list, struct kdnode *item, double dist_sq)
814 struct res_node *rnode;
816 if(!(rnode = alloc_resnode())) {
820 rnode->dist_sq = dist_sq;
823 while(list->next && list->next->dist_sq < dist_sq) {
827 rnode->next = list->next;
832 static void clear_results(struct kdres *rset)
834 struct res_node *tmp, *node = rset->rlist->next;
842 rset->rlist->next = 0;