kdtree_cartesian.cc 13.6 KB
Newer Older
Uwe Schulzweida's avatar
Uwe Schulzweida committed
1
2
3
4
/*!\file kdtree_cartesian.c
 * \brief Routines specific to the n-dimensional cartesian kd-tree
 *
 */
5
6
7
#if defined(HAVE_CONFIG_H)
#include "config.h"
#endif
Uwe Schulzweida's avatar
Uwe Schulzweida committed
8
9
10
11
12
13
14
15
16
17

#include "kdtree.h"
#include "pqueue.h"

/* ********************************************************************

   general utility functions 

   ********************************************************************* */

18
static constexpr
19
kdata_t square(const kdata_t x)
20
21
22
23
{
  return x*x;
}

24
25
static constexpr
kdata_t kd_dist_sq(const kdata_t *restrict a, const kdata_t *restrict b)
26
{
27
  return square((a[0]-b[0]))+square((a[1]-b[1]))+square((a[2]-b[2]));
28
29
}

30

31
kdata_t kd_min(kdata_t x, kdata_t y)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
{
    return x < y ? x : y;
}


/* ******************************************************************
   
   Functions for building and destroying trees 

   ******************************************************************** */


/*! \brief build kd-tree structure
 *
 * \param points an array of kd_points (struct with position vector
 * and data container).
 *
 * \param nPoints the length of the points array.
 *
 * \param constr a pointer to a void *constructor() function to include
 * the data container in the tree; optional, can be NULL
 *
 * \param destr a pointer to a void destructor() function to free()
 * the data containers in the tree; optional, can be NULL, but should
 * be given if the constr argument is non-NULL.
 *
 * \param min a vector with the minimum positions of the corners of the
 * hyperrectangle containing the data.
 *
 * \param max a vector with the maximum positions of the corners of
 * the hyperrectangle containing the data.
 *
 * \param dim the dimensionality of the data.
 *
 * \param max_threads the maximal number of threads spawned for
 * construction of the tree. The threads will be unbalanced if this is
 * not a power of 2.
 *
 * \return root node of the tree
 */
struct kdNode *
73
kd_buildTree(struct kd_point *points, size_t nPoints,
74
             kdata_t *min, kdata_t *max, int dim, int max_threads)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
75
{
76
77
    struct kd_thread_data *my_data = kd_buildArg(points, nPoints, min, max, 0, max_threads, dim);
    struct kdNode *tree = (kdNode *)kd_doBuildTree(my_data);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
    free(my_data);
    return tree;
}


/* *******************************************************************

   Functions for range searches 
  
   *******************************************************************
*/

/* Returns 1 if node is a point in the hyperrectangle defined by
   minimum and maximum vectors min and max. */
int
93
kd_isPointInRect(struct kdNode *node, kdata_t *min, kdata_t *max, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
{
    int i;

    if (node == NULL)
        return 0;

    for(i = 0; i < dim; i++) {
        if (node->location[i] < min[i] || node->location[i] > max[i])
            return 0;
    }
    return 1;
}

/* Returns 1 if the hyperrectangle of node is fully contained within
   the HR described by the minimum and maximum vectors min and
   max. Returns 0 otherwise. */
int
111
kd_isRectInRect(struct kdNode *node, kdata_t *min, kdata_t *max, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
{
    int i;

    if (node == NULL)
        return 0;

    for(i = 0; i < dim; i++) {
        if (node->min[i] < min[i] || node->max[i] > max[i])
            return 0;
    }
    return 1;
}

/* Returns 1 if the hyperrectangle of node overlaps the HR described
   the minimum and maximum vectors min and max. Returns 0
   otherwise. */
int
129
kd_rectOverlapsRect(struct kdNode *node, kdata_t *min, kdata_t *max, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
{
    int i;

    if (node == NULL)
        return 0;

    for(i = 0; i < dim; i++) {
        if ((node->min[i] > max[i] || node->max[i] < min[i]))
            return 0;
    }
    return 1;
}

/*!
 * \brief Perform orthogonal range search (get all points in a
 * hyperrectangle).
 *
 * \param node the root node of tree to be searched.
 *
 * \param min a vector with the minimum positions of the corners of the
 * hyperrectangle containing the data.
 *
 * \param max a vector with the maximum positions of the corners of
 * the hyperrectangle containing the data.
 *
 * \param dim the dimension of the data.
 *
 * \return  Pointer to a priority queue, NULL in case of problems.
*/
struct pqueue *
160
kd_ortRangeSearch(struct kdNode *node, kdata_t *min, kdata_t *max, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
161
162
163
164
165
166
{
    struct pqueue *res;

    if ((res = pqinit(NULL, 1)) == NULL)
        return NULL;
    if (!kd_doOrtRangeSearch(node, min, max, dim, res)) {
167
        for(size_t i = 0; i < res->size; i++) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
168
169
170
171
172
173
174
175
176
177
178
179
            free(res->d[i]);
        }
        free(res->d);
        free(res);
        return NULL;
    }
    return res;
}

/* This is the orthogonal range search. Returns 1 if okay, 0 in case
   of problems. */
int
180
kd_doOrtRangeSearch(struct kdNode *node, kdata_t *min, kdata_t *max,
Uwe Schulzweida's avatar
Uwe Schulzweida committed
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
                    int dim, struct pqueue *res)
{

    if (node == NULL)
        return 1;

    if (kd_isleaf(node) && kd_isPointInRect(node, min, max, dim)) {
        return kd_insertResTree(node, res);
    } else {
        if (kd_isRectInRect(node->left, min, max, dim)) {
            if (!kd_insertResTree(node->left, res))
                return 0;
        } else {
            if (kd_rectOverlapsRect(node->left, min, max, dim))
                if (!kd_doOrtRangeSearch(node->left, min, max, dim, res))
                    return 0;
        }
        if (kd_isRectInRect(node->right, min, max, dim)) {
            if (!kd_insertResTree(node->right, res))
                return 0;
        } else {
            if (kd_rectOverlapsRect(node->right, min, max, dim))
                if (!kd_doOrtRangeSearch(node->right, min, max, dim, res))
                    return 0;
        }
    }
    return 1;
}

/*! 
 * \brief Find the nearest neighbor of a point.
 *
 * \param node the root node of the tree to be searched.
 *
 * \param p a vector to the point whose nearest neighbor is sought.
 *
 * \param max_dist_sq the square of the maximum distance to the
 * nearest neighbor.
 *
 * \param dim the dimension of the data.
 *
 * \return A pointer to node containing the nearest neighbor.
 * max_dist_sq is set to the square of the distance to the nearest
 * neigbor.
 */
struct kdNode *
227
kd_nearest(struct kdNode *node, kdata_t *p, kdata_t *max_dist_sq, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
228
{
229
    if ( !node ) return NULL;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
230

231
232
233
    struct kdNode *nearer, *further;
    if ( p[node->split] < node->location[node->split] ) {
        nearer  = node->left;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
234
235
        further = node->right;
    } else {
236
        nearer  = node->right;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
237
238
        further = node->left;
    }
239
240
241
242
    
    struct kdNode *tmp = kd_nearest(nearer, p, max_dist_sq, dim);
    struct kdNode *nearest = NULL;
    if ( tmp )
Uwe Schulzweida's avatar
Uwe Schulzweida committed
243
244
245
246
        nearest = tmp;
    else
        nearest = node;

247
    kdata_t dist_sq = kd_dist_sq(nearest->location, p);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
248
249
250
251
252
253
    if (*max_dist_sq > dist_sq)
	*max_dist_sq = dist_sq;

    if (!further)
        return nearest;

254
255
256
    kdata_t dx = kd_min(KDATA_ABS(p[node->split] - further->min[node->split]),
                        KDATA_ABS(p[node->split] - further->max[node->split]));
    if ( *max_dist_sq > kd_sqr(dx) ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
257
258
259
260
        /*
         * some part of the further hyper-rectangle is in the search
         * radius, search the further node 
         */
261
        struct kdNode *tmp_nearest;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
262
        tmp = kd_nearest(further, p, max_dist_sq, dim);
263
        if ( tmp )
Uwe Schulzweida's avatar
Uwe Schulzweida committed
264
265
266
267
            tmp_nearest = tmp;
        else
            tmp_nearest = further;

268
        kdata_t tmp_dist_sq = kd_dist_sq(tmp_nearest->location, p);
269
        if ( tmp_dist_sq < dist_sq ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
270
271
272
273
            nearest = tmp_nearest;
            dist_sq = tmp_dist_sq;
            *max_dist_sq = kd_min(dist_sq, *max_dist_sq);
        }
274
        // Uwe Schulzweida
275
        else if ( tmp_dist_sq <= dist_sq && tmp_nearest->index < nearest->index ) {
276
277
278
279
            nearest = tmp_nearest;
            dist_sq = tmp_dist_sq;
            *max_dist_sq = kd_min(dist_sq, *max_dist_sq);
        }
Uwe Schulzweida's avatar
Uwe Schulzweida committed
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
    }
    return nearest;
}

/*! 
 * \brief Return the q nearest-neighbors to a point.
 *
 * \param node the root node of the tree to be searched.
 *
 * \param p a vector to the point whose nearest neighbors are sought.
 *
 * \param max_dist_sq the square of the maximum distance to the
 * nearest neighbors.
 *
 * \param q the maximum number of points to be retured.
 *
 * \param dim the dimension of the data.
 *
 * \return A pointer to a priority queue of the points found, or NULL
 * in case of problems.
 */
struct pqueue *
302
kd_qnearest(struct kdNode *node, kdata_t *p,
303
            kdata_t *max_dist_sq, size_t q, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
304
{
305
306
307
308
    struct pqueue *res = pqinit(NULL, q + 2);
    if ( res == NULL) return NULL;
    
    if ( !kd_doQnearest(node, p, max_dist_sq, q + 1, dim, res) ) {
309
        for ( size_t i = 0; i < res->size; ++i ) free(res->d[i]);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
310
311
312
313
        free(res->d);
        free(res);
        return NULL;
    }
314
    
Uwe Schulzweida's avatar
Uwe Schulzweida committed
315
316
317
318
319
320
321
322
323
324
325
326
    return res;
}

/* *
 * This is the q nearest-neighbor search.
 *
 * This is a modification of the range search in which the maximum
 * search radius is decreased to the maximum of the queue as soon as
 * the queue is filled.
 *
 * return 1 if okay, zero in case of problems
 */
327
328
// Uwe Schulzweida: extract kd_check_dist() from kd_doQnearest()
static int
329
kd_check_dist(struct kdNode *node, kdata_t *p,
330
              kdata_t *max_dist_sq, unsigned int q, struct pqueue *res)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
331
{
332
    kdata_t dist_sq = kd_dist_sq(node->location, p);
333
334
335
    if ( dist_sq < *max_dist_sq && kd_isleaf(node) ) {
        struct resItem *point = (struct resItem *) kd_malloc(sizeof(struct resItem), "kd_doQnearest: ");
        if ( point == NULL) return 0;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
336
337
338
339
        point->node = node;
        point->dist_sq = dist_sq;
        pqinsert(res, point);
    }
340

341
342
    if ( res->size > q ) {
        struct resItem *item;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
343
344
        pqremove_max(res, &item);
        free(item);
345
        if ( res->size > 1 ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
346
347
348
349
350
351
352
353
354
355
356
357
358
            /*
             * Only inspect the queue if there are items left 
             */
            pqpeek_max(res, &item);
            *max_dist_sq = item->dist_sq;
        } else {
            /*
             * Nothing was found within the max search radius 
             */
            *max_dist_sq = 0;
        }
    }

359
360
361
362
    return 1;
}

int
363
kd_doQnearest(struct kdNode *node, kdata_t *p,
364
              kdata_t *max_dist_sq, size_t q, int dim, struct pqueue *res)
365
{
366
    if ( !node ) return 1;
367

368
    if ( !kd_check_dist(node, p, max_dist_sq, q, res) ) return 0;
369

370
371
372
    struct kdNode *nearer, *further;
    if ( p[node->split] < node->location[node->split] ) {
        nearer  = node->left;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
373
374
        further = node->right;
    } else {
375
        nearer  = node->right;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
376
377
        further = node->left;
    }
378
    if ( !kd_doQnearest(nearer, p, max_dist_sq, q, dim, res) ) return 0;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
379

380
    if ( !further ) return 1;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
381

382
383
384
385
    kdata_t dx = kd_min(KDATA_ABS(p[node->split] - further->min[node->split]),
                        KDATA_ABS(p[node->split] - further->max[node->split]));
    
    if ( *max_dist_sq > kd_sqr(dx) ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
386
387
388
389
        /*
         * some part of the further hyper-rectangle is in the search
         * radius, search the further node 
         */
390
        if (!kd_doQnearest(further, p, max_dist_sq, q, dim, res)) return 0;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
391

392
        if (!kd_check_dist(node, p, max_dist_sq, q, res)) return 0;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
    }
    return 1;
}


/*!
 * \brief Perform a range search around a point.
 *
 * \param node the root node of the tree to be searched.
 *
 * \param p the location of the point around which the search is carried out .
 *
 * \param max_dist_sq the square of the radius of the hypersphere.
 *
 * \param dim the dimension of the data.  \param ordered determines
 * whether the result list should be ordered in increasing distance
 * (KD_ORDERED) or unordered (KD_UNORDERED).
 *
 * \return A pointer to a priority queue containing the points found,
 * NULL in case of problems.
 */
struct pqueue *
415
kd_range(struct kdNode *node, kdata_t *p, kdata_t *max_dist_sq,
Uwe Schulzweida's avatar
Uwe Schulzweida committed
416
417
         int dim, int ordered)
{
418
419
420
421
    struct pqueue *res = pqinit(NULL, 1);
    if ( res == NULL ) return NULL;
    
    if ( !kd_doRange(node, p, max_dist_sq, dim, res, ordered) ) {
422
        for( size_t i = 0; i < res->size; ++i ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
423
424
425
426
427
428
429
430
431
432
433
434
            free(res->d[i]);
        }
        free(res->d);
        free(res);
        return NULL;
    }
    return res;
}


/* This is the range search. Returns 1 if okay, 0 in case of problems */
int
435
kd_doRange(struct kdNode *node, kdata_t *p, kdata_t *max_dist_sq,
Uwe Schulzweida's avatar
Uwe Schulzweida committed
436
437
438
439
440
           int dim, struct pqueue *res, int ordered)
{

    struct kdNode *nearer, *further;
    struct resItem *point;
441
    kdata_t dist_sq, dx;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
442
443
444
445

    if (!node)
        return 1;

446
    dist_sq = kd_dist_sq(node->location, p);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
447
    if (dist_sq < *max_dist_sq && kd_isleaf(node)) {
Oliver Heidmann's avatar
Oliver Heidmann committed
448
        if ((point = (struct resItem *)kd_malloc(sizeof(struct resItem), "kd_doRange:"))
Uwe Schulzweida's avatar
Uwe Schulzweida committed
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
            == NULL)
            return 0;
        point->node = node;
        point->dist_sq = ordered ? dist_sq : -1;
        pqinsert(res, point);
    }

    if (p[node->split] < node->location[node->split]) {
        nearer = node->left;
        further = node->right;
    } else {
        nearer = node->right;
        further = node->left;
    }
    if (!kd_doRange(nearer, p, max_dist_sq, dim, res, ordered))
        return 0;

    if (!further)
        return 1;

469
470
    dx = kd_min(KDATA_ABS(p[node->split] - further->min[node->split]),
                KDATA_ABS(p[node->split] - further->max[node->split]));
Uwe Schulzweida's avatar
Uwe Schulzweida committed
471
472
473
474
475
476
477
    if (*max_dist_sq > kd_sqr(dx)) {
        /*
         * some part of the further hyper-rectangle is in the search
         * radius, search the further node 
         */
        if (!kd_doRange(further, p, max_dist_sq, dim, res, ordered))
            return 0;
478
        dist_sq = kd_dist_sq(node->location, p);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
479
480

        if (dist_sq < *max_dist_sq && kd_isleaf(node)) {
Oliver Heidmann's avatar
Oliver Heidmann committed
481
            if ((point = (struct resItem *)kd_malloc(sizeof(struct resItem), "kd_doRange: "))
Uwe Schulzweida's avatar
Uwe Schulzweida committed
482
483
484
485
486
487
488
489
490
491
492
493
                == NULL)
                return 0;
            point->node = node;
            point->dist_sq = ordered ? dist_sq : -1;
            pqinsert(res, point);
        }

    }
    return 1;
}

/* End range searching functions */