41 #include "pf_vector.h" 42 #include "pf_kdtree.h" 46 static int pf_kdtree_equal(pf_kdtree_t *
self,
int key_a[],
int key_b[]);
49 static pf_kdtree_node_t *pf_kdtree_insert_node(pf_kdtree_t *
self, pf_kdtree_node_t *parent,
50 pf_kdtree_node_t *node,
int key[],
double value);
53 static pf_kdtree_node_t *pf_kdtree_find_node(pf_kdtree_t *
self, pf_kdtree_node_t *node,
int key[]);
56 static void pf_kdtree_cluster_node(pf_kdtree_t *
self, pf_kdtree_node_t *node,
int depth);
65 static void pf_kdtree_draw_node(pf_kdtree_t *
self, pf_kdtree_node_t *node, rtk_fig_t *fig);
73 pf_kdtree_t *pf_kdtree_alloc(
int max_size)
77 self = calloc(1,
sizeof(pf_kdtree_t));
81 self->size[2] = (10 * M_PI / 180);
86 self->node_max_count = max_size;
87 self->nodes = calloc(self->node_max_count,
sizeof(pf_kdtree_node_t));
97 void pf_kdtree_free(pf_kdtree_t *
self)
107 void pf_kdtree_clear(pf_kdtree_t *
self)
110 self->leaf_count = 0;
111 self->node_count = 0;
119 void pf_kdtree_insert(pf_kdtree_t *
self, pf_vector_t pose,
double value)
123 key[0] = floor(pose.v[0] / self->size[0]);
124 key[1] = floor(pose.v[1] / self->size[1]);
125 key[2] = floor(pose.v[2] / self->size[2]);
127 self->root = pf_kdtree_insert_node(
self, NULL, self->root, key, value);
157 double pf_kdtree_get_prob(pf_kdtree_t *
self, pf_vector_t pose)
160 pf_kdtree_node_t *node;
162 key[0] = floor(pose.v[0] / self->size[0]);
163 key[1] = floor(pose.v[1] / self->size[1]);
164 key[2] = floor(pose.v[2] / self->size[2]);
166 node = pf_kdtree_find_node(
self, self->root, key);
175 int pf_kdtree_get_cluster(pf_kdtree_t *
self, pf_vector_t pose)
178 pf_kdtree_node_t *node;
180 key[0] = floor(pose.v[0] / self->size[0]);
181 key[1] = floor(pose.v[1] / self->size[1]);
182 key[2] = floor(pose.v[2] / self->size[2]);
184 node = pf_kdtree_find_node(
self, self->root, key);
187 return node->cluster;
193 int pf_kdtree_equal(pf_kdtree_t *
self,
int key_a[],
int key_b[])
197 if (key_a[0] != key_b[0])
199 if (key_a[1] != key_b[1])
202 if (key_a[2] != key_b[2])
222 pf_kdtree_node_t *pf_kdtree_insert_node(pf_kdtree_t *
self, pf_kdtree_node_t *parent,
223 pf_kdtree_node_t *node,
int key[],
double value)
226 int split, max_split;
231 assert(self->node_count < self->node_max_count);
232 node =
self->nodes +
self->node_count++;
233 memset(node, 0,
sizeof(pf_kdtree_node_t));
240 node->depth = parent->depth + 1;
242 for (i = 0; i < 3; i++)
243 node->key[i] = key[i];
246 self->leaf_count += 1;
253 if (pf_kdtree_equal(
self, key, node->key))
255 node->value += value;
264 node->pivot_dim = -1;
265 for (i = 0; i < 3; i++)
267 split = abs(key[i] - node->key[i]);
268 if (split > max_split)
274 assert(node->pivot_dim >= 0);
276 node->pivot_value = (key[node->pivot_dim] + node->key[node->pivot_dim]) / 2.0;
278 if (key[node->pivot_dim] < node->pivot_value)
280 node->children[0] = pf_kdtree_insert_node(
self, node, NULL, key, value);
281 node->children[1] = pf_kdtree_insert_node(
self, node, NULL, node->key, node->value);
285 node->children[0] = pf_kdtree_insert_node(
self, node, NULL, node->key, node->value);
286 node->children[1] = pf_kdtree_insert_node(
self, node, NULL, key, value);
290 self->leaf_count -= 1;
297 assert(node->children[0] != NULL);
298 assert(node->children[1] != NULL);
300 if (key[node->pivot_dim] < node->pivot_value)
301 pf_kdtree_insert_node(
self, node, node->children[0], key, value);
303 pf_kdtree_insert_node(
self, node, node->children[1], key, value);
312 pf_kdtree_node_t *pf_kdtree_find_node(pf_kdtree_t *
self, pf_kdtree_node_t *node,
int key[])
319 if (pf_kdtree_equal(
self, key, node->key))
328 assert(node->children[0] != NULL);
329 assert(node->children[1] != NULL);
332 if (key[node->pivot_dim] < node->pivot_value)
333 return pf_kdtree_find_node(
self, node->children[0], key);
335 return pf_kdtree_find_node(
self, node->children[1], key);
365 void pf_kdtree_cluster(pf_kdtree_t *
self)
368 int queue_count, cluster_count;
369 pf_kdtree_node_t **queue, *node;
372 queue = calloc(self->node_count,
sizeof(queue[0]));
375 for (i = 0; i <
self->node_count; i++)
377 node =
self->nodes + i;
381 assert(queue_count < self->node_count);
382 queue[queue_count++] = node;
385 assert(node == pf_kdtree_find_node(
self, self->root, node->key));
392 while (queue_count > 0)
394 node = queue[--queue_count];
397 if (node->cluster >= 0)
401 node->cluster = cluster_count++;
404 pf_kdtree_cluster_node(
self, node, 0);
414 void pf_kdtree_cluster_node(pf_kdtree_t *
self, pf_kdtree_node_t *node,
int depth)
418 pf_kdtree_node_t *nnode;
420 for (i = 0; i < 3 * 3 * 3; i++)
422 nkey[0] = node->key[0] + (i / 9) - 1;
423 nkey[1] = node->key[1] + ((i % 9) / 3) - 1;
424 nkey[2] = node->key[2] + ((i % 9) % 3) - 1;
426 nnode = pf_kdtree_find_node(
self, self->root, nkey);
434 if (nnode->cluster >= 0)
436 assert(nnode->cluster == node->cluster);
441 nnode->cluster = node->cluster;
443 pf_kdtree_cluster_node(
self, nnode, depth + 1);
450 #ifdef INCLUDE_RTKGUI 454 void pf_kdtree_draw(pf_kdtree_t *
self, rtk_fig_t *fig)
456 if (self->root != NULL)
457 pf_kdtree_draw_node(
self, self->root, fig);
464 void pf_kdtree_draw_node(pf_kdtree_t *
self, pf_kdtree_node_t *node, rtk_fig_t *fig)
471 ox = (node->key[0] + 0.5) *
self->size[0];
472 oy = (node->key[1] + 0.5) *
self->size[1];
474 rtk_fig_rectangle(fig, ox, oy, 0.0, self->size[0], self->size[1], 0);
479 snprintf(text,
sizeof(text),
"%d", node->cluster);
480 rtk_fig_text(fig, ox, oy, 0.0, text);
484 assert(node->children[0] != NULL);
485 assert(node->children[1] != NULL);
486 pf_kdtree_draw_node(
self, node->children[0], fig);
487 pf_kdtree_draw_node(
self, node->children[1], fig);