#include "kdtree.h"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
enum { DIM = 3 };
struct kdhyperrect {
float min[DIM], max[DIM];
};
struct kdnode {
float pos[DIM];
int dir;
void *data;
struct kdnode *left, *right;
};
struct kdtree {
struct kdnode *root;
struct kdhyperrect *rect;
void (*destr)(void *);
};
#define SQ(x) ((x) * (x))
static void clear_rec(struct kdnode *node, void (*destr)(void *));
static int insert_rec(struct kdnode **node, const float *pos, void *data,
int dir);
static struct kdhyperrect *hyperrect_create(const float *min, const float *max);
static void hyperrect_extend(struct kdhyperrect *rect, const float *pos);
static float hyperrect_dist_sq(struct kdhyperrect *rect, const float *pos);
struct kdtree *kd_create() {
struct kdtree *tree;
if (!(tree = malloc(sizeof *tree))) {
return 0;
}
tree->root = 0;
tree->destr = 0;
tree->rect = 0;
return tree;
}
static void clear_rec(struct kdnode *node, void (*destr)(void *)) {
if (!node)
return;
clear_rec(node->left, destr);
clear_rec(node->right, destr);
if (destr) {
destr(node->data);
}
free(node);
}
static void kd_clear(struct kdtree *tree) {
clear_rec(tree->root, tree->destr);
tree->root = 0;
if (tree->rect) {
free(tree->rect);
tree->rect = 0;
}
}
void kd_free(struct kdtree *tree) {
if (tree) {
kd_clear(tree);
free(tree);
}
}
static int insert_rec(struct kdnode **nptr, const float *pos, void *data,
int dir) {
int new_dir;
struct kdnode *node;
if (!*nptr) {
if (!(node = malloc(sizeof *node))) {
return -1;
}
memcpy(node->pos, pos, DIM * sizeof *node->pos);
node->data = data;
node->dir = dir;
node->left = node->right = 0;
*nptr = node;
return 0;
}
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);
}
return insert_rec(&(*nptr)->right, pos, data, new_dir);
}
int kd_insert(struct kdtree *tree, const float *pos, void *data) {
if (insert_rec(&tree->root, pos, data, 0)) {
return -1;
}
if (tree->rect == 0) {
tree->rect = hyperrect_create(pos, pos);
} else {
hyperrect_extend(tree->rect, pos);
}
return 0;
}
static void kd_nearest_i(struct kdnode *node, const float *pos,
struct kdnode **result, float *result_dist_sq,
struct kdhyperrect *rect) {
int dir = node->dir;
int i;
float dummy, dist_sq;
struct kdnode *nearer_subtree, *farther_subtree;
float *nearer_hyperrect_coord, *farther_hyperrect_coord;
dummy = pos[dir] - node->pos[dir];
if (dummy <= 0) {
nearer_subtree = node->left;
farther_subtree = node->right;
nearer_hyperrect_coord = rect->max + dir;
farther_hyperrect_coord = rect->min + dir;
} else {
nearer_subtree = node->right;
farther_subtree = node->left;
nearer_hyperrect_coord = rect->min + dir;
farther_hyperrect_coord = rect->max + dir;
}
if (nearer_subtree) {
dummy = *nearer_hyperrect_coord;
*nearer_hyperrect_coord = node->pos[dir];
kd_nearest_i(nearer_subtree, pos, result, result_dist_sq, rect);
*nearer_hyperrect_coord = dummy;
}
dist_sq = 0;
for (i = 0; i < DIM; i++) {
dist_sq += SQ(node->pos[i] - pos[i]);
}
if (dist_sq < *result_dist_sq) {
*result = node;
*result_dist_sq = dist_sq;
}
if (farther_subtree) {
dummy = *farther_hyperrect_coord;
*farther_hyperrect_coord = node->pos[dir];
if (hyperrect_dist_sq(rect, pos) < *result_dist_sq) {
kd_nearest_i(farther_subtree, pos, result, result_dist_sq, rect);
}
*farther_hyperrect_coord = dummy;
}
}
int kd_nearest(struct kdtree *kd, const float *pos, void **res) {
struct kdhyperrect rect;
struct kdnode *result;
float dist_sq;
int i;
if (!kd)
return -1;
if (!kd->rect)
return -1;
memcpy(&rect, kd->rect, sizeof(struct kdhyperrect));
result = kd->root;
dist_sq = 0;
for (i = 0; i < DIM; i++)
dist_sq += SQ(result->pos[i] - pos[i]);
kd_nearest_i(kd->root, pos, &result, &dist_sq, &rect);
if (result) {
*res = result->data;
return 0;
} else {
return -1;
}
}
static struct kdhyperrect *hyperrect_create(const float *min,
const float *max) {
size_t size = DIM * sizeof(float);
struct kdhyperrect *rect = 0;
if (!(rect = malloc(sizeof(struct kdhyperrect)))) {
return 0;
}
memcpy(rect->min, min, size);
memcpy(rect->max, max, size);
return rect;
}
static void hyperrect_extend(struct kdhyperrect *rect, const float *pos) {
int i;
for (i = 0; i < DIM; i++) {
if (pos[i] < rect->min[i]) {
rect->min[i] = pos[i];
}
if (pos[i] > rect->max[i]) {
rect->max[i] = pos[i];
}
}
}
static float hyperrect_dist_sq(struct kdhyperrect *rect, const float *pos) {
int i;
float result = 0;
for (i = 0; i < DIM; i++) {
if (pos[i] < rect->min[i]) {
result += SQ(rect->min[i] - pos[i]);
} else if (pos[i] > rect->max[i]) {
result += SQ(rect->max[i] - pos[i]);
}
}
return result;
}