Implementing the kd_nearest_n function

examples/Makefile

@@ -5,11 +5,12 @@ CFLAGS = -std=c89 -pedantic -Wall -g -I..
 LDFLAGS = $(kdlib) -lm
 
 .PHONY: all
-all: test test2
+all: test test2 test3
 
 test: test.c $(LDFLAGS)
 test2: test2.c $(LDFLAGS)
+test3: test3.c $(LDFLAGS)
 
 .PHONY: clean
 clean:
-	rm -f test test2
+	rm -f test test2 test3

examples/test3.c

@@ -0,0 +1,33 @@
+#include <stdio.h>
+#include "kdtree.h"
+
+
+int main(void) {
+    void* tree = kd_create(2);
+    void* res;
+    double data[] = {
+        2, 3,
+        5, 4,
+        9, 6,
+        4, 7,
+        8, 1,
+        7, 2
+    };
+    double pos[2] = {9, 2};
+    kd_insert(tree, &data[0], NULL);
+    kd_insert(tree, &data[2], NULL);
+    kd_insert(tree, &data[4], NULL);
+    kd_insert(tree, &data[6], NULL);
+    kd_insert(tree, &data[8], NULL);
+    kd_insert(tree, &data[10], NULL);
+
+    res = kd_nearest_n(tree, pos, 3);
+    while(!kd_res_end(res)) {
+        kd_res_item(res, &data[0]);
+        printf("%f %f %f\n", data[0], data[1], kd_res_dist(res));
+        kd_res_next(res);
+    }
+    kd_res_free(res);
+	kd_free(tree);
+    return 0;
+}

kdtree.c

@@ -87,10 +87,10 @@ struct kdres {
 
 #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 int rlist_insert(struct res_node *list, struct kdnode *item, double dist_sq);
+static struct res_node *rlist_pop_back(struct res_node *list);
 static void clear_results(struct kdres *set);
 
 static struct kdhyperrect* hyperrect_create(int dim, const double *min, const double *max);
@@ -102,6 +102,7 @@ static double hyperrect_dist_sq(struct kdhyperrect *rect, const double *pos);
 #ifdef USE_LIST_NODE_ALLOCATOR
 static struct res_node *alloc_resnode(void);
 static void free_resnode(struct res_node*);
+static void free_resnode_buffer();
 #else
 #define alloc_resnode()		malloc(sizeof(struct res_node))
 #define free_resnode(n)		free(n)
@@ -131,6 +132,10 @@ void kd_free(struct kdtree *tree)
 		kd_clear(tree);
 		free(tree);
 	}
+
+#ifdef USE_LIST_NODE_ALLOCATOR
+	free_resnode_buffer();
+#endif
 }
 
 static void clear_rec(struct kdnode *node, void (*destr)(void*))
@@ -158,7 +163,7 @@ void kd_clear(struct kdtree *tree)
 	}
 }
 
-void kd_data_destructor(struct kdtree *tree, void (*destr)(void*))
+void kd_data_destructor(struct kdtree *tree, void (*destr)(void *))
 {
 	tree->destr = destr;
 }
@@ -292,55 +297,41 @@ static int find_nearest(struct kdnode *node, const double *pos, double range, st
 	return added_res;
 }
 
-#if 0
-static int find_nearest_n(struct kdnode *node, const double *pos, double range, int num, struct rheap *heap, int dim)
+static int find_nearest_n(struct kdnode *node, const double *pos, int num, int *size, double *dist_max, struct res_node *list, int dim)
 {
 	double dist_sq, dx;
-	int i, ret, added_res = 0;
+	int i, ret;
 
 	if(!node) return 0;
-
-	/* if the photon is close enough, add it to the result heap */
+
 	dist_sq = 0;
 	for(i=0; i<dim; i++) {
 		dist_sq += SQ(node->pos[i] - pos[i]);
 	}
-	if(dist_sq <= range_sq) {
-		if(heap->size >= num) {
-			/* get furthest element */
-			struct res_node *maxelem = rheap_get_max(heap);
-
-			/* and check if the new one is closer than that */
-			if(maxelem->dist_sq > dist_sq) {
-				rheap_remove_max(heap);
-
-				if(rheap_insert(heap, node, dist_sq) == -1) {
-					return -1;
-				}
-				added_res = 1;
 
-				range_sq = dist_sq;
-			}
+	if(dist_sq < *dist_max) {
+		if(*size < num) {
+			++(*size);
 		} else {
-			if(rheap_insert(heap, node, dist_sq) == -1) {
-				return =1;
-			}
-			added_res = 1;
+			struct res_node *back = rlist_pop_back(list);
+			if (back == 0)
+				return -1;
+			*dist_max = back->dist_sq > dist_sq ? back->dist_sq : dist_sq;
+		}
+		if (rlist_insert(list, node, dist_sq) == -1) {
+			return -1;
 		}
 	}
 
-
 	/* find signed distance from the splitting plane */
 	dx = pos[node->dir] - node->pos[node->dir];
 
-	ret = find_nearest_n(dx <= 0.0 ? node->left : node->right, pos, range, num, heap, dim);
-	if(ret >= 0 && fabs(dx) < range) {
-		added_res += ret;
-		ret = find_nearest_n(dx <= 0.0 ? node->right : node->left, pos, range, num, heap, dim);
+	ret = find_nearest_n(dx <= 0.0 ? node->left : node->right, pos, num, size, dist_max, list, dim);
+	if(ret >= 0 && SQ(dx) < *dist_max) {
+		ret = find_nearest_n(dx <= 0.0 ? node->right : node->left, pos, num, size, dist_max, list, dim);
 	}
-
+	return ret;
 }
-#endif
 
 static void kd_nearest_i(struct kdnode *node, const double *pos, struct kdnode **result, double *result_dist_sq, struct kdhyperrect* rect)
 {
@@ -509,11 +500,12 @@ struct kdres *kd_nearest3f(struct kdtree *tree, float x, float y, float z)
 }
 
 /* ---- nearest N search ---- */
-/*
-static kdres *kd_nearest_n(struct kdtree *kd, const double *pos, int num)
+
+struct kdres *kd_nearest_n(struct kdtree *kd, const double *pos, int num)
 {
-	int ret;
+	int ret, size = 0;
 	struct kdres *rset;
+	double dist_max = DBL_MAX;
 
 	if(!(rset = malloc(sizeof *rset))) {
 		return 0;
@@ -525,14 +517,67 @@ static kdres *kd_nearest_n(struct kdtree *kd, const double *pos, int num)
 	rset->rlist->next = 0;
 	rset->tree = kd;
 
-	if((ret = find_nearest_n(kd->root, pos, range, num, rset->rlist, kd->dim)) == -1) {
+	ret = find_nearest_n(kd->root, pos, num, &size, &dist_max, rset->rlist, kd->dim);
+	if(ret == -1) {
 		kd_res_free(rset);
 		return 0;
 	}
-	rset->size = ret;
+	rset->size = size;
 	kd_res_rewind(rset);
 	return rset;
-}*/
+}
+
+struct kdres *kd_nearest_nf(struct kdtree *tree, const float *pos, int num)
+{
+	static double sbuf[16];
+	double *bptr, *buf = 0;
+	int dim = tree->dim;
+	struct kdres *res;
+
+	if (dim > 16) {
+#ifndef NO_ALLOCA
+		if (dim <= 256)
+			bptr = buf = alloca(dim * sizeof *bptr);
+		else
+#endif
+		    if (!(bptr = buf = malloc(dim * sizeof *bptr))) {
+			return 0;
+		}
+	} else {
+		bptr = buf = sbuf;
+	}
+
+	while (dim-- > 0) {
+		*bptr++ = *pos++;
+	}
+
+	res = kd_nearest_n(tree, buf, num);
+#ifndef NO_ALLOCA
+	if (tree->dim > 256)
+#else
+	if (tree->dim > 16)
+#endif
+		free(buf);
+	return res;
+}
+
+struct kdres *kd_nearest_n3(struct kdtree *tree, double x, double y, double z, int num)
+{
+	double pos[3];
+	pos[0] = x;
+	pos[1] = y;
+	pos[2] = z;
+	return kd_nearest_n(tree, pos, num);
+}
+
+struct kdres *kd_nearest_n3f(struct kdtree *tree, float x, float y, float z, int num)
+{
+	double pos[3];
+	pos[0] = x;
+	pos[1] = y;
+	pos[2] = z;
+	return kd_nearest_n(tree, pos, num);
+}
 
 struct kdres *kd_nearest_range(struct kdtree *kd, const double *pos, double range)
 {
@@ -690,6 +735,11 @@ void *kd_res_item_data(struct kdres *set)
 	return kd_res_item(set, 0);
 }
 
+double kd_res_dist(struct kdres *set)
+{
+	return sqrt(set->riter->dist_sq);
+}
+
 /* ---- hyperrectangle helpers ---- */
 static struct kdhyperrect* hyperrect_create(int dim, const double *min, const double *max)
 {
@@ -804,8 +854,29 @@ static void free_resnode(struct res_node *node)
 	pthread_mutex_unlock(&alloc_mutex);
 #endif
 }
-#endif	/* list node allocator or not */
 
+static void free_resnode_buffer()
+{
+#ifndef NO_PTHREADS
+	pthread_mutex_lock(&alloc_mutex);
+#endif
+
+	if (free_nodes) {
+		struct res_node *ptr = free_nodes;
+		while (ptr) {
+			ptr = ptr->next;
+			free(free_nodes);
+			free_nodes = ptr;
+		}
+		free_nodes = 0;
+	}
+
+#ifndef NO_PTHREADS
+	pthread_mutex_unlock(&alloc_mutex);
+#endif
+}
+
+#endif /* list node allocator or not */
 
 /* inserts the item. if dist_sq is >= 0, then do an ordered insert */
 /* TODO make the ordering code use heapsort */
@@ -829,6 +900,20 @@ static int rlist_insert(struct res_node *list, struct kdnode *item, double dist_
 	return 0;
 }
 
+static struct res_node *rlist_pop_back(struct res_node *list)
+{
+	struct res_node *previous = 0;
+	while (list->next) {
+		previous = list;
+		list = list->next;
+	}
+	if (previous) {
+		previous->next = 0;
+	}
+	free_resnode(list);
+	return previous;
+}
+
 static void clear_results(struct kdres *rset)
 {
 	struct res_node *tmp, *node = rset->rlist->next;
@@ -840,4 +925,4 @@ static void clear_results(struct kdres *rset)
 	}
 
 	rset->rlist->next = 0;
-}
+}

kdtree.h

@@ -73,12 +73,10 @@ struct kdres *kd_nearest3f(struct kdtree *tree, float x, float y, float z);
  * 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_nearest_n(struct kdtree *tree, const double *pos, int num);
 struct kdres *kd_nearest_nf(struct kdtree *tree, const float *pos, int num);
-struct kdres *kd_nearest_n3(struct kdtree *tree, double x, double y, double z);
-struct kdres *kd_nearest_n3f(struct kdtree *tree, float x, float y, float z);
-*/
+struct kdres *kd_nearest_n3(struct kdtree *tree, double x, double y, double z, int num);
+struct kdres *kd_nearest_n3f(struct kdtree *tree, float x, float y, float z, int num);
 
 /* Find any nearest nodes from a given point within a range.
  *
@@ -121,6 +119,9 @@ void *kd_res_item3f(struct kdres *set, float *x, float *y, float *z);
 /* equivalent to kd_res_item(set, 0) */
 void *kd_res_item_data(struct kdres *set);
 
+/* returns the distance between the requested position and the found point.
+ */
+double kd_res_dist(struct kdres *set);
 
 #ifdef __cplusplus
 }