KDTree.cpp

Go to the documentation of this file.

namespace krotos
{
 
KDTree::KDTree()
{
    //
}
 
KDTree::~KDTree()
{
    //
}
 
void KDTree::addDatasetItem(float x, float y, float z)
{
    std::vector<float> pt = {x, y, z};
    dataset.push_back(pt);
}
 
void KDTree::buildIndex() { kdTree = std::make_unique<myKDTree>(3, dataset, 10); }
 
std::size_t KDTree::knnQuery(float x, float y, float z)
{
    assert(kdTree);
 
    const std::size_t num_results = 1;
    std::vector<float> query_pt = {x, y, z};
 
    // do a knn search
    std::vector<std::size_t> ret_indexes(num_results);
    std::vector<float> out_dists_sqr(num_results);
 
    nanoflann::KNNResultSet<float> resultSet(num_results);
 
    resultSet.init(&ret_indexes[0], &out_dists_sqr[0]);
    kdTree->index->findNeighbors(resultSet, &query_pt[0]);
    return ret_indexes[0];
}
 
std::vector<std::size_t> KDTree::knnQuery(float x, float y, float z, std::size_t k)
{
    assert(kdTree);
 
    std::vector<float> query_pt = {x, y, z};
 
    // do a knn search
    std::vector<std::size_t> ret_indexes(k);
    std::vector<float> out_dists_sqr(k);
 
    nanoflann::KNNResultSet<float> resultSet(k);
 
    resultSet.init(&ret_indexes[0], &out_dists_sqr[0]);
    kdTree->index->findNeighbors(resultSet, &query_pt[0]);
 
    return ret_indexes;
}
 
std::vector<std::size_t> KDTree::radiusSearch(float x, float y, float z, float search_radius)
{
    assert(kdTree);
 
    std::vector<float> query_pt = {x, y, z};
 
    std::vector<nanoflann::ResultItem<std::size_t, float>> matches;
 
    nanoflann::SearchParameters params;
    params.sorted = false;
 
    const std::size_t nMatches = kdTree->index->radiusSearch(&query_pt[0], search_radius, matches, params);
 
    std::vector<std::size_t> ret_indexes(nMatches);
    for (std::size_t i = 0; i < nMatches; i++)
        ret_indexes[i] = matches[i].first;
    return ret_indexes;
}
 
// void KDTree::removePoint(std::size_t idx) {
//    kdTree->index->removePoint(idx);
//}
 
} // namespace krotos