KDTree_fwd.hpp

Go to the documentation of this file.

#ifndef MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_FWD_HPP_
#define MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_FWD_HPP_
 
#include <medusa/Config.hpp>
#include <nanoflann/nanoflann.hpp>
#include <medusa/bits/types/Range_fwd.hpp>
#include <iosfwd>
#include <medusa/bits/utils/memutils.hpp>
#include "PointCloud.hpp"
 
namespace mm {
 
template <class vec_t>
class KDTree {
  public:
    typedef vec_t vector_t;   
    typedef typename vector_t::scalar_t scalar_t;   
    enum {  dim = vec_t::dim };
 
  private:
    typedef nanoflann::KDTreeSingleIndexAdaptor<
            nanoflann::L2_Simple_Adaptor<scalar_t, kdtree_internal::PointCloud<vec_t>>,
            kdtree_internal::PointCloud<vec_t>, dim, int> kd_tree_t;   
 
    kdtree_internal::PointCloud<vec_t> points_;  
    kd_tree_t tree;  
 
  public:
    explicit KDTree(const Range<vec_t>& points) :
            points_(points), tree(dim, points_, nanoflann::KDTreeSingleIndexAdaptorParams(20)) {
        tree.buildIndex();
    }
 
    KDTree() : points_(), tree(dim, points_, nanoflann::KDTreeSingleIndexAdaptorParams(20)) {}
 
    void reset(const Range<vec_t>& points) {
        points_.setPts(points);
        tree.buildIndex();
    }
 
    std::pair<Range<int>, Range<scalar_t>> query(const vec_t& point, int k = 1) const;
 
    std::pair<Range<int>, Range<scalar_t>> query(
            const vec_t& point, const scalar_t& radius_squared) const;
 
    vec_t get(int index) const { return points_.get(index); }
 
    Range<vec_t> get(const Range<int>& indexes) const {
        const int n = indexes.size();
        Range<vec_t> result(n);
        for (int i = 0; i < n; ++i) {
            result[i] = points_.get(indexes[i]);
        }
        return result;
    }
 
    int size() const { return points_.kdtree_get_point_count(); }
 
    template <class V>
    friend std::ostream& operator<<(std::ostream& os, const KDTree<V>& tree);
};
 
}  // namespace mm
 
#endif  // MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_FWD_HPP_