From 8e6ced70cc48dc842b23eaed5c60fb72ae266661 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Tue, 11 Mar 2014 12:42:06 -0400 Subject: Initial commit. --- kd-forest.c | 302 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 302 insertions(+) create mode 100644 kd-forest.c (limited to 'kd-forest.c') diff --git a/kd-forest.c b/kd-forest.c new file mode 100644 index 0000000..4a5bf37 --- /dev/null +++ b/kd-forest.c @@ -0,0 +1,302 @@ +/********************************************************************* + * kd-forest * + * Copyright (C) 2014 Tavian Barnes * + * * + * 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 +#include +#include + +void +kd_node_init(kd_node_t *node, unsigned int x, unsigned int y) +{ + node->left = node->right = NULL; + node->x = x; + node->y = y; + node->added = node->removed = false; +} + +void +kd_node_set_color(kd_node_t *node, uint32_t color) +{ + node->coords[0] = (color >> 16) & 0xFF; + node->coords[1] = (color >> 8) & 0xFF; + node->coords[2] = color & 0xFF; +} + +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); + } + return count; +} + +typedef int kd_comparator(const void *a, const void* b); + +static int kd_compare0(const void *a, const void *b) { + int aval = (*(const kd_node_t **)a)->coords[0]; + int bval = (*(const kd_node_t **)b)->coords[0]; + return aval - bval; +} + +static int kd_compare1(const void *a, const void *b) { + int aval = (*(const kd_node_t **)a)->coords[1]; + int bval = (*(const kd_node_t **)b)->coords[1]; + return aval - bval; +} + +static int kd_compare2(const void *a, const void *b) { + int aval = (*(const kd_node_t **)a)->coords[2]; + int bval = (*(const kd_node_t **)b)->coords[2]; + return 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 int +kd_distance_sq(kd_node_t *a, kd_node_t *b) +{ + int result = 0; + for (int i = 0; i < KD_DIMEN; ++i) { + int d = a->coords[i] - b->coords[i]; + result += d*d; + } + return result; +} + +static void +kd_find_nearest_recursive(kd_node_t *root, kd_node_t *target, kd_node_t **best, + int *limit, unsigned int coord) +{ + int dist = target->coords[coord] - root->coords[coord]; + int dist_sq = dist*dist; + + if (!root->removed) { + int root_dist_sq = kd_distance_sq(root, target); + if (!*best || root_dist_sq < *limit) { + *best = root; + *limit = root_dist_sq; + } + } + + coord = (coord + 1)%KD_DIMEN; + + if (root->left && (dist <= 0 || !*best || dist_sq <= *limit)) { + kd_find_nearest_recursive(root->left, target, best, limit, coord); + } + if (root->right && (dist >= 0 || !*best || dist_sq <= *limit)) { + kd_find_nearest_recursive(root->right, target, best, limit, coord); + } +} + +static void +kd_find_nearest(kd_node_t *root, kd_node_t *target, kd_node_t **best, + int *limit) +{ + kd_find_nearest_recursive(root, 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) +{ + 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) +{ + node->added = true; + + // 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(kd_forest_t *kdf, kd_node_t *target) +{ + int limit; + 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; +} -- cgit v1.2.3