kdtree.hpp

// SPDX-License-Identifier: MIT
// Copyright (c) 2021  Kenji Koide (k.koide@aist.go.jp)
 
#pragma once
 
#include <memory>
#include <iostream>
#include <Eigen/Core>
 
#include <gtsam_points/ann/nearest_neighbor_search.hpp>
 
namespace gtsam_points {
 
template <typename PointCloud, typename Projection>
struct UnsafeKdTree;
struct AxisAlignedProjection;
 
struct KdTree : public NearestNeighborSearch {
public:
  using Index = UnsafeKdTree<KdTree, AxisAlignedProjection>;
 
  KdTree(const Eigen::Vector4d* points, int num_points, int build_num_threads = 1);
  virtual ~KdTree() override;
 
  virtual size_t knn_search(
    const double* pt,
    size_t k,
    size_t* k_indices,
    double* k_sq_dists,
    double max_sq_dist = std::numeric_limits<double>::max()) const override;
 
  virtual size_t radius_search(
    const double* pt,
    double radius,
    std::vector<size_t>& indices,
    std::vector<double>& sq_dists,
    int max_num_neighbors = std::numeric_limits<int>::max()) const override;
 
public:
  const int num_points;
  const Eigen::Vector4d* points;
 
  double search_eps;
 
  std::unique_ptr<Index> index;
};
 
}  // namespace gtsam_points