You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

52 lines
1.9 KiB

#ifndef MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_HPP_
#define MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_HPP_
/**
* @file
* Implementation of KDTree.
*/
#include "KDTree_fwd.hpp"
#include <medusa/bits/utils/assert.hpp>
#include <medusa/bits/utils/memutils.hpp>
namespace mm {
template <class vec_t>
std::pair<Range<int>, Range<typename vec_t::scalar_t>> KDTree<vec_t>::query(const vec_t& point,
int k) const {
assert_msg(point.array().isFinite().prod() == 1, "Invalid point.");
Range<int> ret_index(k);
Range<scalar_t> out_dist_sqr(k);
int actual_k = tree.knnSearch(point.data(), k, &ret_index[0], &out_dist_sqr[0]);
assert_msg(actual_k == k, "There were not enough points in the tree, you requested %d "
"points, the tree only contains %d points.", k, actual_k);
return {ret_index, out_dist_sqr};
}
template <class vec_t>
std::pair<Range<int>, Range<typename vec_t::scalar_t>> KDTree<vec_t>::query(
const vec_t& point, const scalar_t& radius_squared) const {
assert_msg(point.array().isFinite().prod() == 1, "Invalid point.");
std::vector<std::pair<int, scalar_t>> idx_dist;
int k = tree.radiusSearch(point.data(), radius_squared, idx_dist, nanoflann::SearchParams());
Range<int> idx(k); Range<scalar_t> dists(k);
for (int i = 0; i < k; ++i) {
std::tie(idx[i], dists[i]) = idx_dist[i];
}
return {idx, dists};
}
/// Output basic info about given tree.
template <class V>
std::ostream& operator<<(std::ostream& os, const KDTree<V>& tree) {
return os << "KDTree:\n"
<< " dimension: " << tree.dim << '\n'
<< " num of points: " << tree.size() << "\n"
<< " first point: " << tree.get(0) << "\n"
<< " last point: " << tree.get(tree.size() - 1) << "\n";
}
} // namespace mm
#endif // MEDUSA_BITS_SPATIAL_SEARCH_KDTREE_HPP_