74 lines
1.5 KiB
C++
74 lines
1.5 KiB
C++
#ifndef AKNN_H
|
|
#define AKNN_H
|
|
|
|
#include "nanoflann.hpp"
|
|
|
|
using namespace nanoflann;
|
|
|
|
template <typename T, int dim> class aKNN {
|
|
|
|
struct DataSet {
|
|
|
|
std::vector<T> elems;
|
|
|
|
inline size_t kdtree_get_point_count() const {return elems.size();}
|
|
inline float kdtree_distance(const float* p1, const size_t idxP2, size_t) const {
|
|
float dist = 0;
|
|
for (int i = 0; i < dim; ++i) {
|
|
float delta = (p1[i] - kdtree_get_pt(idxP2, i));
|
|
dist += delta*delta;
|
|
}
|
|
return dist;
|
|
}
|
|
|
|
inline float kdtree_get_pt(const size_t idx, int pos) const {
|
|
return elems[idx].feature[pos];
|
|
}
|
|
|
|
template <class BBOX> bool kdtree_get_bbox(BBOX&) const {return false;}
|
|
|
|
} data;
|
|
|
|
typedef KDTreeSingleIndexAdaptor<L2_Simple_Adaptor<float, DataSet>, DataSet, dim> MyTree;
|
|
|
|
MyTree* tree = nullptr;
|
|
|
|
|
|
|
|
public:
|
|
|
|
/** add a new element */
|
|
void add(const T& elem) {
|
|
data.elems.push_back(elem);
|
|
}
|
|
|
|
/** build the KD-Tree */
|
|
void build() {
|
|
tree = new MyTree(dim, data, KDTreeSingleIndexAdaptorParams(10) );
|
|
tree->buildIndex();
|
|
}
|
|
|
|
/** get the nearest n elements */
|
|
template <typename T2> std::vector<T> get(const T2* query, const int numResults) {
|
|
|
|
float distances[numResults];
|
|
size_t indices[numResults];
|
|
|
|
KNNResultSet<float> res(numResults);
|
|
res.init(indices, distances);
|
|
|
|
tree->knnSearch(query, numResults, indices, distances);
|
|
std::vector<T> vec;
|
|
for (int i = 0; i < numResults; ++i) {
|
|
vec.push_back(data.elems[indices[i]]);
|
|
}
|
|
return vec;
|
|
|
|
}
|
|
|
|
|
|
};
|
|
|
|
#endif // AKNN_H
|
|
|