/** * Find n nearest returns number of points found, with results in nearest. * Normal is optional, but if given will limit results to points in normal direction from co. * * \param r_nearest An array of nearest, sized at least \a n. */ int BLI_kdtree_find_nearest_n__normal( KDTree *tree, const float co[3], const float nor[3], KDTreeNearest r_nearest[], unsigned int n) { KDTreeNode *root, *node = NULL; KDTreeNode **stack, *defaultstack[KD_STACK_INIT]; float cur_dist; unsigned int totstack, cur = 0; unsigned int i, found = 0; #ifdef DEBUG BLI_assert(tree->is_balanced == true); #endif if (!tree->root || n == 0) return 0; stack = defaultstack; totstack = KD_STACK_INIT; root = tree->root; cur_dist = squared_distance(root->co, co, nor); add_nearest(r_nearest, &found, n, root->index, cur_dist, root->co); if (co[root->d] < root->co[root->d]) { if (root->right) stack[cur++] = root->right; if (root->left) stack[cur++] = root->left; } else { if (root->left) stack[cur++] = root->left; if (root->right) stack[cur++] = root->right; } while (cur--) { node = stack[cur]; cur_dist = node->co[node->d] - co[node->d]; if (cur_dist < 0.0f) { cur_dist = -cur_dist * cur_dist; if (found < n || -cur_dist < r_nearest[found - 1].dist) { cur_dist = squared_distance(node->co, co, nor); if (found < n || cur_dist < r_nearest[found - 1].dist) add_nearest(r_nearest, &found, n, node->index, cur_dist, node->co); if (node->left) stack[cur++] = node->left; } if (node->right) stack[cur++] = node->right; } else { cur_dist = cur_dist * cur_dist; if (found < n || cur_dist < r_nearest[found - 1].dist) { cur_dist = squared_distance(node->co, co, nor); if (found < n || cur_dist < r_nearest[found - 1].dist) add_nearest(r_nearest, &found, n, node->index, cur_dist, node->co); if (node->right) stack[cur++] = node->right; } if (node->left) stack[cur++] = node->left; } if (UNLIKELY(cur + 3 > totstack)) { stack = realloc_nodes(stack, &totstack, defaultstack != stack); } } for (i = 0; i < found; i++) r_nearest[i].dist = sqrtf(r_nearest[i].dist); if (stack != defaultstack) MEM_freeN(stack); return (int)found; }
/* finds the nearest n entries in tree to specified coordinates */ int BLI_kdtree_find_n_nearest(KDTree *tree, int n, const float co[3], const float nor[3], KDTreeNearest *nearest) { KDTreeNode *root, *node = NULL; KDTreeNode **stack, *defaultstack[100]; float cur_dist; int i, totstack, cur = 0, found = 0; if (!tree->root) return 0; stack = defaultstack; totstack = 100; root = tree->root; cur_dist = squared_distance(root->co, co, root->nor, nor); add_nearest(nearest, &found, n, root->index, cur_dist, root->co); if (co[root->d] < root->co[root->d]) { if (root->right) stack[cur++] = root->right; if (root->left) stack[cur++] = root->left; } else { if (root->left) stack[cur++] = root->left; if (root->right) stack[cur++] = root->right; } while (cur--) { node = stack[cur]; cur_dist = node->co[node->d] - co[node->d]; if (cur_dist < 0.0f) { cur_dist = -cur_dist * cur_dist; if (found < n || -cur_dist < nearest[found - 1].dist) { cur_dist = squared_distance(node->co, co, node->nor, nor); if (found < n || cur_dist < nearest[found - 1].dist) add_nearest(nearest, &found, n, node->index, cur_dist, node->co); if (node->left) stack[cur++] = node->left; } if (node->right) stack[cur++] = node->right; } else { cur_dist = cur_dist * cur_dist; if (found < n || cur_dist < nearest[found - 1].dist) { cur_dist = squared_distance(node->co, co, node->nor, nor); if (found < n || cur_dist < nearest[found - 1].dist) add_nearest(nearest, &found, n, node->index, cur_dist, node->co); if (node->right) stack[cur++] = node->right; } if (node->left) stack[cur++] = node->left; } if (cur + 3 > totstack) { KDTreeNode **temp = MEM_callocN((totstack + 100) * sizeof(KDTreeNode *), "psys_treestack"); memcpy(temp, stack, totstack * sizeof(KDTreeNode *)); if (stack != defaultstack) MEM_freeN(stack); stack = temp; totstack += 100; } } for (i = 0; i < found; i++) nearest[i].dist = sqrt(nearest[i].dist); if (stack != defaultstack) MEM_freeN(stack); return found; }