knn_result.hpp

// SPDX-FileCopyrightText: Copyright 2024 Kenji Koide
// SPDX-License-Identifier: MIT
#pragma once
 
#include <limits>
#include <iostream>
 
namespace gtsam_points {
 
struct KnnSetting {
public:
  template <typename Result>
  bool fulfilled(const Result& result) const {
    return result.worst_distance() < epsilon || result.num_found() >= max_nn;
  }
 
public:
  int max_nn = std::numeric_limits<int>::max();             
  double max_sq_dist = std::numeric_limits<double>::max();  
  double epsilon = 0.0;                                     
};
 
struct identity_transform {
  template <typename T>
  T operator()(const T& x) const {
    return x;
  }
};
 
template <int N, typename IndexTransform = identity_transform>
struct KnnResult {
public:
  static constexpr size_t INVALID = std::numeric_limits<size_t>::max();
 
  explicit KnnResult(
    size_t* indices,
    double* distances,
    int num_neighbors = -1,
    const IndexTransform& index_transform = IndexTransform(),
    double max_sq_dist = std::numeric_limits<double>::max())
  : index_transform(index_transform),
    capacity(num_neighbors),
    num_found_neighbors(0),
    indices(indices),
    distances(distances) {
    if constexpr (N > 0) {
      if (num_neighbors >= 0) {
        std::cerr << "warning: Specifying dynamic num_neighbors=" << num_neighbors << " for a static KNN result container (N=" << N << ")"
                  << std::endl;
        abort();
      }
    } else {
      if (num_neighbors <= 0) {
        std::cerr << "error: Specifying invalid num_neighbors=" << num_neighbors << " for a dynamic KNN result container" << std::endl;
        abort();
      }
    }
 
    std::fill(this->indices, this->indices + buffer_size(), INVALID);
    std::fill(this->distances, this->distances + buffer_size(), max_sq_dist);
  }
 
  size_t buffer_size() const {
    if constexpr (N > 0) {
      return N;
    } else {
      return capacity;
    }
  }
 
  size_t num_found() const { return num_found_neighbors; }
 
  double worst_distance() const { return distances[buffer_size() - 1]; }
 
  void push(size_t index, double distance) {
    if (distance >= worst_distance()) {
      return;
    }
 
    if constexpr (N == 1) {
      indices[0] = index_transform(index);
      distances[0] = distance;
    } else {
      int insert_loc = std::min<int>(num_found_neighbors, buffer_size() - 1);
      for (; insert_loc > 0 && distance < distances[insert_loc - 1]; insert_loc--) {
        indices[insert_loc] = indices[insert_loc - 1];
        distances[insert_loc] = distances[insert_loc - 1];
      }
 
      indices[insert_loc] = index_transform(index);
      distances[insert_loc] = distance;
    }
 
    num_found_neighbors = std::min<int>(num_found_neighbors + 1, buffer_size());
  }
 
public:
  const IndexTransform& index_transform;
  const int capacity;       
  int num_found_neighbors;  
  size_t* indices;          
  double* distances;        
};
 
template <typename IndexTransform = identity_transform>
struct RadiusSearchResult {
public:
  explicit RadiusSearchResult(const IndexTransform& index_transform = IndexTransform()) : index_transform(index_transform) { neighbors.reserve(32); }
 
  size_t num_found() const { return neighbors.size(); }
 
  constexpr double worst_distance() const { return std::numeric_limits<double>::max(); }
 
  void sort() {
    std::sort(neighbors.begin(), neighbors.end(), [](const auto& a, const auto& b) { return a.second < b.second; });
  }
 
  void push(size_t index, double distance) { neighbors.emplace_back(index_transform(index), distance); }
 
public:
  const IndexTransform& index_transform;
  std::vector<std::pair<size_t, double>> neighbors;  
};
 
}  // namespace gtsam_points