Skip to content
Permalink
d95e93bf70
Switch branches/tags
Go to file
 
 
Cannot retrieve contributors at this time
327 lines (280 sloc) 8.07 KB
/*********************************************************************
* kd-forest *
* Copyright (C) 2014 Tavian Barnes <tavianator@tavianator.com> *
* *
* This program is free software. It comes without any warranty, to *
* the extent permitted by applicable law. You can redistribute it *
* and/or modify it under the terms of the Do What The Fuck You Want *
* To Public License, Version 2, as published by Sam Hocevar. See *
* the COPYING file or http://www.wtfpl.net/ for more details. *
*********************************************************************/
#include "kd-forest.h"
#include "util.h"
#include <math.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
kd_node_t *
new_kd_node(double coords[KD_DIMEN], unsigned int x, unsigned int y)
{
kd_node_t *node = xmalloc(sizeof(kd_node_t));
memcpy(node->coords, coords, sizeof(node->coords));
node->left = node->right = NULL;
node->x = x;
node->y = y;
node->removed = false;
return node;
}
static void
kd_destroy(kd_node_t *node)
{
if (node) {
kd_destroy(node->left);
kd_destroy(node->right);
free(node);
}
}
static size_t
kd_collect_nodes(kd_node_t *root, kd_node_t **buffer, bool include_removed)
{
size_t count = 0;
if (include_removed || !root->removed) {
buffer[0] = root;
++count;
}
if (root->left) {
count += kd_collect_nodes(root->left, buffer + count, include_removed);
}
if (root->right) {
count += kd_collect_nodes(root->right, buffer + count, include_removed);
}
if (root->removed && !include_removed) {
free(root);
}
return count;
}
typedef int kd_comparator(const void *a, const void* b);
static int kd_compare0(const void *a, const void *b) {
double aval = (*(const kd_node_t **)a)->coords[0];
double bval = (*(const kd_node_t **)b)->coords[0];
return (aval > bval) - (aval < bval);
}
static int kd_compare1(const void *a, const void *b) {
double aval = (*(const kd_node_t **)a)->coords[1];
double bval = (*(const kd_node_t **)b)->coords[1];
return (aval > bval) - (aval < bval);
}
static int kd_compare2(const void *a, const void *b) {
double aval = (*(const kd_node_t **)a)->coords[2];
double bval = (*(const kd_node_t **)b)->coords[2];
return (aval > bval) - (aval < bval);
}
static kd_comparator *kd_comparators[KD_DIMEN] = {
kd_compare0,
kd_compare1,
kd_compare2,
};
// When building k-d trees, we use KD_DIMEN sorted arrays of nodes plus one
// extra array for scratch space
#define KD_BUFSIZE (KD_DIMEN + 1)
static kd_node_t *
kd_build_tree_recursive(kd_node_t **buffers[KD_BUFSIZE], size_t size, unsigned int coord)
{
if (size == 0) {
return NULL;
}
size_t split = size/2;
size_t left_size = split, right_size = size - left_size - 1;
kd_node_t *root = buffers[coord][split];
for (size_t i = 0; i < size; ++i) {
buffers[coord][i]->is_left = i < left_size;
}
kd_node_t **right_buffers[KD_BUFSIZE];
for (int i = 0; i < KD_DIMEN; ++i) {
right_buffers[i] = buffers[i] + left_size + 1;
}
kd_node_t **scratch = buffers[KD_DIMEN];
right_buffers[KD_DIMEN] = scratch;
for (size_t i = 0; i < KD_DIMEN; ++i) {
if (i == coord) {
continue;
}
kd_node_t **buffer = buffers[i];
kd_node_t **right_buffer = right_buffers[i];
for (size_t j = 0, k = 0, skip = 0; j < size; ++j) {
if (buffer[j]->is_left) {
buffer[j - skip] = buffer[j];
} else {
if (buffer[j] != root) {
scratch[k] = buffer[j];
++k;
}
++skip;
}
}
for (size_t j = 0; j < right_size; ++j) {
right_buffer[j] = scratch[j];
}
}
coord = (coord + 1)%KD_DIMEN;
root->left = kd_build_tree_recursive(buffers, left_size, coord);
root->right = kd_build_tree_recursive(right_buffers, right_size, coord);
return root;
}
static kd_node_t *
kd_build_tree(kd_node_t **buffers[KD_BUFSIZE], size_t size)
{
for (int i = 1; i < KD_DIMEN; ++i) {
memcpy(buffers[i], buffers[0], size*sizeof(kd_node_t *));
}
for (int i = 0; i < KD_DIMEN; ++i) {
qsort(buffers[i], size, sizeof(kd_node_t *), kd_comparators[i]);
}
return kd_build_tree_recursive(buffers, size, 0);
}
static double
kd_distance_sq(const double a[KD_DIMEN], const double b[KD_DIMEN])
{
double result = 0.0;
for (int i = 0; i < KD_DIMEN; ++i) {
double d = a[i] - b[i];
result += d*d;
}
return result;
}
static void
kd_find_nearest_recursive(kd_node_t *node, const double target[KD_DIMEN], const double closest[KD_DIMEN], kd_node_t **best, double *limit, unsigned int coord)
{
if (!node->removed) {
double node_dist_sq = kd_distance_sq(node->coords, target);
if (node_dist_sq < *limit) {
*limit = node_dist_sq;
*best = node;
}
}
kd_node_t *first;
kd_node_t *second;
if (target[coord] < node->coords[coord]) {
first = node->left;
second = node->right;
} else {
first = node->right;
second = node->left;
}
unsigned int next = (coord + 1)%KD_DIMEN;
if (first) {
kd_find_nearest_recursive(first, target, closest, best, limit, next);
}
if (second) {
double new_closest[KD_DIMEN];
memcpy(new_closest, closest, sizeof(new_closest));
new_closest[coord] = node->coords[coord];
if (kd_distance_sq(new_closest, target) < *limit) {
kd_find_nearest_recursive(second, target, new_closest, best, limit, next);
}
}
}
static void
kd_find_nearest(kd_node_t *root, const double target[KD_DIMEN], kd_node_t **best, double *limit)
{
kd_find_nearest_recursive(root, target, target, best, limit, 0);
}
void
kdf_init(kd_forest_t *kdf)
{
kdf->roots = NULL;
kdf->size = kdf->size_est = 0;
kdf->roots_size = kdf->roots_capacity = 0;
}
void
kdf_destroy(kd_forest_t *kdf)
{
for (unsigned int i = 0; i < kdf->roots_size; ++i) {
kd_destroy(kdf->roots[i]);
}
free(kdf->roots);
}
static size_t
kdf_collect_nodes(kd_forest_t *kdf, kd_node_t **buffer, unsigned int slot, bool include_removed)
{
size_t count = 0;
for (unsigned int i = 0; i < slot; ++i) {
if (kdf->roots[i]) {
count += kd_collect_nodes(kdf->roots[i], buffer + count, include_removed);
}
}
return count;
}
static void
kdf_balance(kd_forest_t *kdf, kd_node_t *new_node, bool force)
{
++kdf->size;
size_t slot, buffer_size;
if (force) {
buffer_size = kdf->size_est = kdf->size;
slot = kdf->roots_size;
} else {
++kdf->size_est;
for (slot = 0; slot < kdf->roots_size; ++slot) {
if (!kdf->roots[slot]) {
break;
}
}
buffer_size = 1 << slot;
}
kd_node_t **buffer = xmalloc(buffer_size*sizeof(kd_node_t *));
buffer[0] = new_node;
kdf_collect_nodes(kdf, buffer + 1, slot, !force);
kd_node_t **buffers[KD_BUFSIZE];
for (int i = 1; i < KD_BUFSIZE; ++i) {
buffers[i] = xmalloc(buffer_size*sizeof(kd_node_t *));
}
if (slot >= kdf->roots_capacity) {
kdf->roots_capacity = slot + 1;
kdf->roots = xrealloc(kdf->roots, kdf->roots_capacity*sizeof(kd_node_t *));
}
size_t i, offset;
for (i = 0, offset = 0; offset < buffer_size; ++i) {
size_t chunk_size = 1 << i;
if (buffer_size & chunk_size) {
buffers[0] = buffer + offset;
kdf->roots[i] = kd_build_tree(buffers, chunk_size);
offset |= chunk_size;
} else {
kdf->roots[i] = NULL;
}
}
if (force || i > kdf->roots_size) {
kdf->roots_size = i;
}
free(buffer);
for (i = 1; i < KD_BUFSIZE; ++i) {
free(buffers[i]);
}
}
void
kdf_insert(kd_forest_t *kdf, kd_node_t *node)
{
// If half or more of the nodes are removed, force a complete rebalance
bool force = (kdf->size_est + 1) >= 2*(kdf->size + 1);
kdf_balance(kdf, node, force);
}
void
kdf_remove(kd_forest_t *kdf, kd_node_t *node)
{
--kdf->size;
node->removed = true;
}
kd_node_t *
kdf_find_nearest(const kd_forest_t *kdf, const double target[KD_DIMEN])
{
double limit = INFINITY;
kd_node_t *best = NULL;
for (unsigned int i = 0; i < kdf->roots_size; ++i) {
kd_node_t *root = kdf->roots[i];
if (root != NULL) {
kd_find_nearest(root, target, &best, &limit);
}
}
return best;
}