Index: trunk/psLib/src/types/psTree.c =================================================================== --- trunk/psLib/src/types/psTree.c (revision 18312) +++ trunk/psLib/src/types/psTree.c (revision 18317) @@ -294,5 +294,9 @@ } -static void treePrint(FILE *fptr, const psTreeNode *node, int level) +// Print a representation of the node; called recursively +static void treeNodePrint(FILE *fptr, // File pointer to which to print + const psTreeNode *node, // Node to print + int level // Level at which to print + ) { long index = node->index; // Index of node @@ -322,6 +326,6 @@ if (node->left && node->right) { fprintf(fptr, " ==> %lf\n", tree->division[index]); - treePrint(fptr, node->left, level); - treePrint(fptr, node->right, level); + treeNodePrint(fptr, node->left, level); + treeNodePrint(fptr, node->right, level); return; } @@ -354,5 +358,5 @@ PS_ASSERT_TREE_NON_NULL(tree, ); - treePrint(fptr, tree->root, 0); + treeNodePrint(fptr, tree->root, 0); return; } @@ -375,5 +379,8 @@ // Returns the square of the distance between a point on the tree and an aribtary point -static inline double treeContentDistance(const psTree *tree, long index, const psVector *coords) +static inline double treeContentDistance(const psTree *tree, // Tree of interest + long index, // Index of point on tree + const psVector *coords // Coordinates of interest + ) { int dim = tree->dim; // Dimensionality @@ -418,5 +425,8 @@ // Returns whether the designated leaf contains an arbitrary point -static inline bool treeBranchInside(const psTree *tree, long index, const psVector *coords) +static inline bool treeBranchInside(const psTree *tree, // Tree of interest + long index, // Index of leaf on tree + const psVector *coords // Coordinates of interest + ) { int dim = tree->dim; // Dimensionality @@ -433,10 +443,10 @@ // Search a leaf node for a closer point -static inline void treeLeafSearch(long *bestIndex, - double *bestDistance, - const psTree *tree, - const psTreeNode *leaf, - const psVector *coords - ) +static inline void treeLeafSearchNearest(long *bestIndex, // Index of closest point + double *bestDistance, // Distance to closest point + const psTree *tree, // Tree of interest + const psTreeNode *leaf, // Leaf to search + const psVector *coords // Coordinates of interest + ) { for (int i = 0; i < leaf->num; i++) { @@ -452,5 +462,4 @@ - // Given an arbitrary point, return the index of the nearest neighbour long psTreeNearest(const psTree *tree, const psVector *coords) @@ -467,5 +476,5 @@ long bestIndex = -1; // Index of best point double bestDistance = INFINITY; // Distance to best point - treeLeafSearch(&bestIndex, &bestDistance, tree, leaf, coords); + treeLeafSearchNearest(&bestIndex, &bestDistance, tree, leaf, coords); // Work up from the leaf containing the point of interest, and for each in turn, search down the other @@ -494,5 +503,5 @@ } // Leaf node - treeLeafSearch(&bestIndex, &bestDistance, tree, node, coords); + treeLeafSearchNearest(&bestIndex, &bestDistance, tree, node, coords); work->data[workIndex] = NULL; workIndex--; @@ -504,4 +513,137 @@ return bestIndex; +} + +// Search a leaf node for points within radius squared +static inline long treeLeafSearchWithin(double radius2, // Radius squared to search + const psTree *tree, // Tree of interest + const psTreeNode *leaf, // Leaf to search + const psVector *coords // Coordinates of interest + ) +{ + long num = 0; // Number within circle + for (int i = 0; i < leaf->num; i++) { + long index = leaf->contents[i]; // Index of point + double distance = treeContentDistance(tree, index, coords); // Distance to point + if (distance < radius2) { + num++; + } + } + return num; +} + +// Given an arbitrary point and a matching radius, return the number of points within that radius +long psTreeWithin(const psTree *tree, const psVector *coords, double radius) +{ +#if 1 // Might be in a tight loop + PS_ASSERT_TREE_NON_NULL(tree, -1); + PS_ASSERT_VECTOR_NON_NULL(coords, -1); + PS_ASSERT_VECTOR_TYPE(coords, PS_TYPE_F64, -1); + PS_ASSERT_VECTOR_SIZE(coords, (long)tree->dim, -1); +#endif + + radius *= radius; // We work with the radius-squared + long num = 0; // Number of points in circle + + // This is essentially the same as psTreeNearest, except we're not allowed to prune + + // Find the closest point in the leaf that contains the point of interest + psTreeNode *leaf = psTreeLeaf(tree, coords); // Leaf containing the point of interest + num += treeLeafSearchWithin(radius, tree, leaf, coords); + + psArray *work = psArrayAlloc(tree->numNodes); // Work queue + while (leaf->up) { + psTreeNode *up = leaf->up; // Parent node + + long workIndex = 0; + work->data[workIndex] = (leaf == up->left ? up->right : up->left); // The other node + while (workIndex >= 0) { + psTreeNode *node = work->data[workIndex]; + if (node->left) { + // Branch node + work->data[workIndex] = node->right; + work->data[++workIndex] = node->left; + continue; + } + // Leaf node + num += treeLeafSearchWithin(radius, tree, node, coords); + work->data[workIndex] = NULL; + workIndex--; + } + + leaf = up; + } + psFree(work); + + return num; +} + +// Search a leaf node for any points within radius squared +static inline bool treeLeafSearchWithinAny(double radius2, // Radius squared to search + const psTree *tree, // Tree of interest + const psTreeNode *leaf, // Leaf to search + const psVector *coords // Coordinates of interest + ) +{ + for (int i = 0; i < leaf->num; i++) { + long index = leaf->contents[i]; // Index of point + double distance = treeContentDistance(tree, index, coords); // Distance to point + if (distance < radius2) { + return true; + } + } + return false; +} + +// Given an arbitrary point and a matching radius, return whether there are any points within that radius +bool psTreeWithinAny(const psTree *tree, const psVector *coords, double radius) +{ +#if 1 // Might be in a tight loop + PS_ASSERT_TREE_NON_NULL(tree, false); + PS_ASSERT_VECTOR_NON_NULL(coords, false); + PS_ASSERT_VECTOR_TYPE(coords, PS_TYPE_F64, false); + PS_ASSERT_VECTOR_SIZE(coords, (long)tree->dim, false); +#endif + + radius *= radius; // We work with the radius-squared + + // This is essentially the same as psTreeWithin, except we can bail as soon as we find something + + // Find the closest point in the leaf that contains the point of interest + psTreeNode *leaf = psTreeLeaf(tree, coords); // Leaf containing the point of interest + if (treeLeafSearchWithinAny(radius, tree, leaf, coords)) { + return true; + } + + psArray *work = psArrayAlloc(tree->numNodes); // Work queue + while (leaf->up) { + psTreeNode *up = leaf->up; // Parent node + + long workIndex = 0; + work->data[workIndex] = (leaf == up->left ? up->right : up->left); // The other node + while (workIndex >= 0) { + psTreeNode *node = work->data[workIndex]; + if (node->left) { + // Branch node + work->data[workIndex] = node->right; + work->data[++workIndex] = node->left; + continue; + } + // Leaf node + if (treeLeafSearchWithinAny(radius, tree, node, coords)) { + // Clear out the work queue + memset(work->data, 0, (workIndex + 1) * sizeof(void)); + psFree(work); + return true; + } + work->data[workIndex] = NULL; + workIndex--; + } + + leaf = up; + } + psFree(work); + + return false; } Index: trunk/psLib/src/types/psTree.h =================================================================== --- trunk/psLib/src/types/psTree.h (revision 18312) +++ trunk/psLib/src/types/psTree.h (revision 18317) @@ -89,4 +89,16 @@ ); +/// Return the number of points within some radius of given coordinates +long psTreeWithin(const psTree *tree, ///< Tree + const psVector *coords, ///< Coordinates of interest + double radius ///< Radius of interest + ); + +/// Return whether there are any points within some radius of given coordinates +bool psTreeWithinAny(const psTree *tree, ///< Tree + const psVector *coords, ///< Coordinates of interest + double radius ///< Radius of interest + ); + /// Return the coordinates of a point in the tree, specified by its index psVector *psTreeCoords(psVector *out, ///< Output vector, or NULL