FrankE
e6329e1db4
deleting from the grid is now bleaching fast added new helper methods many new test-cases many new methods for geo classes and others added a bunch of new grid-walkers
114 lines
3.2 KiB
C++
114 lines
3.2 KiB
C++
#ifndef KNN_H
|
|
#define KNN_H
|
|
|
|
#include "../lib/nanoflann/nanoflann.hpp"
|
|
#include "Debug.h"
|
|
|
|
/**
|
|
* helper class to extract k-nearest-neighbors
|
|
* from a given input data structure.
|
|
* uses nanoflann
|
|
*
|
|
* usage:
|
|
* Grid<30, T> theGrid;
|
|
* KNN<Grid<20, T>, 3, float> knn(theGrid);
|
|
* std::vector<T> elems = knn.get({0,0,0}, 10);
|
|
*/
|
|
template <typename DataStructure, int dim, typename Scalar = float> class KNN {
|
|
|
|
private:
|
|
|
|
static constexpr const char* name = "KNN";
|
|
|
|
/** type-definition for the nanoflann KD-Tree used for searching */
|
|
typedef nanoflann::KDTreeSingleIndexAdaptor<nanoflann::L2_Simple_Adaptor<Scalar, DataStructure>, DataStructure, dim> Tree;
|
|
|
|
/** the maximum depth of the tree */
|
|
static constexpr int maxLeafs = 10;
|
|
|
|
/** the constructed tree used for searching */
|
|
Tree tree;
|
|
|
|
/** the underlying data-structure we want to search within */
|
|
DataStructure& data;
|
|
|
|
public:
|
|
|
|
/** ctor */
|
|
KNN(DataStructure& data) : tree(dim, data, nanoflann::KDTreeSingleIndexAdaptorParams(maxLeafs)), data(data) {
|
|
|
|
Log::add(name, "building kd-tree for " + std::to_string(data.kdtree_get_point_count()) + " elements", false);
|
|
Log::tick();
|
|
tree.buildIndex();
|
|
Log::tock();
|
|
|
|
}
|
|
|
|
/** get the k-nearest-neighbors for the given input point */
|
|
template <typename Element> std::vector<Element> get(const Scalar* point, const int numNeighbors, const float maxDistSquared = 99999) const {
|
|
|
|
// buffer for to-be-fetched neighbors
|
|
size_t indices[numNeighbors];
|
|
float distances[numNeighbors];
|
|
|
|
// find k-nearest-neighbors
|
|
tree.knnSearch(point, numNeighbors, indices, distances);
|
|
|
|
// construct output
|
|
std::vector<Element> elements;
|
|
for (int i = 0; i < numNeighbors; ++i) {
|
|
if (distances[i] > maxDistSquared) {continue;} // too far away?
|
|
elements.push_back(data[indices[i]]);
|
|
}
|
|
return elements;
|
|
|
|
}
|
|
|
|
/** get the k-nearest-neighbors for the given input point */
|
|
template <typename Element> std::vector<Element> get(std::initializer_list<Scalar> point, const int numNeighbors, const float maxDistSquared = 99999) const {
|
|
return get(point.begin(), numNeighbors, maxDistSquared);
|
|
}
|
|
|
|
/** get the nearest neighbor and its distance */
|
|
void getNearest(const Scalar* point, size_t& idx, float& distSquared) {
|
|
|
|
// find 1-nearest-neighbors
|
|
tree.knnSearch(point, 1, &idx, &distSquared);
|
|
|
|
}
|
|
|
|
/** get the index of the element nearest to the given point */
|
|
size_t getNearestIndex(const Scalar* point) {
|
|
size_t idx;
|
|
float distSquared;
|
|
tree.knnSearch(point, 1, &idx, &distSquared);
|
|
return idx;
|
|
}
|
|
|
|
/** get the index of the element nearest to the given point */
|
|
size_t getNearestIndex(const std::initializer_list<Scalar> lst) {
|
|
size_t idx;
|
|
float distSquared;
|
|
tree.knnSearch(lst.begin(), 1, &idx, &distSquared);
|
|
return idx;
|
|
}
|
|
|
|
/** get the distance to the element nearest to the given point */
|
|
float getNearestDistance(const std::initializer_list<Scalar> lst) {
|
|
size_t idx;
|
|
float distSquared;
|
|
tree.knnSearch(lst.begin(), 1, &idx, &distSquared);
|
|
return std::sqrt(distSquared);
|
|
}
|
|
|
|
void get(const Scalar* point, const int numNeighbors, size_t* indices, float* squaredDist) {
|
|
|
|
// find k-nearest-neighbors
|
|
tree.knnSearch(point, numNeighbors, indices, squaredDist);
|
|
|
|
}
|
|
|
|
};
|
|
|
|
#endif // KNN_H
|