NearestNeighbourSearch2D.cpp

Go to the documentation of this file.

#include "NearestNeighbourSearch2D.h"
 
namespace krotos
{
 
NearestNeighbourSearch2D::NearestNeighbourSearch2D(int max_elements)
{
    // dataset.data.reserve(dataset.dim * max_elements);
 
    // m_spaceInstance = std::make_unique <hnswlib::L2Space>(dataset.dim);
    // The line below is for compatability with hnswlib. space is a ptr to hnswlib::SpaceInterface,
    // an abstract class which can't be instantiated as is - it is used as an interface class
    // Assigning m_spaceInstance raw pointer here should be safe because m_spaceInstance
    // will persist until owning class is destroyed
    // If anyone can see a better way to do this outside of re-factoring hnswlib, please do !
    space = m_spaceInstance.get();
    alg = std::make_unique<hnswlib::BruteforceSearch<float>>(space, 2 * max_elements);
}
 
NearestNeighbourSearch2D::~NearestNeighbourSearch2D()
{
    alg.reset();
    m_spaceInstance.reset();
}
 
void NearestNeighbourSearch2D::addDatasetItem(float x, float y)
{
    dataset.data.push_back(x);
    dataset.data.push_back(y);
    alg->addPoint(dataset.data.data() + dataset.index * dataset.dim, dataset.index);
    ++dataset.index;
}
 
int NearestNeighbourSearch2D::knnQuery(float x, float y)
{
    assert(!dataset.data.empty());
 
    std::array<float, 2> features = {x, y};
    auto result = alg->searchKnn(features.data(), 1);
    return static_cast<int>(result.top().second);
}
 
std::vector<int> NearestNeighbourSearch2D::knnQuery(float x, float y, int k)
{
    assert(k <= dataset.index);
 
    std::array<float, 2> features = {x, y};
 
    // avoid k larger than dataset
    k = k <= dataset.index ? k : dataset.index;
 
    std::vector<int> indices;
    auto results = alg->searchKnnCloserFirst(features.data(), k);
    for (const auto& result : results)
        indices.push_back(static_cast<int>(result.second));
 
    return indices;
}
} // namespace krotos