31 #ifndef RTABMAP_FLANN_KDTREE_INDEX_H_
32 #define RTABMAP_FLANN_KDTREE_INDEX_H_
61 (*this)[
"trees"] = trees;
72 template <
typename Distance>
73 class KDTreeIndex :
public NNIndex<Distance>
114 template<
typename Archive>
123 bool leaf_node =
false;
124 if (Archive::is_saving::value) {
130 if (Archive::is_loading::value) {
136 if (Archive::is_loading::value) {
216 size_t old_size =
size_;
223 for (
size_t i=old_size;
i<
size_;++
i) {
237 template<
typename Archive>
246 if (Archive::is_loading::value) {
250 if (Archive::is_loading::value) {
256 if (Archive::is_loading::value) {
265 serialization::SaveArchive sa(
stream);
297 int maxChecks = searchParams.checks;
298 float epsError = 1+searchParams.eps;
302 getExactNeighbors<true>(
result, vec, epsError);
305 getExactNeighbors<false>(
result, vec, epsError);
310 getNeighbors<true>(
result, vec, maxChecks, epsError);
313 getNeighbors<false>(result, vec, maxChecks, epsError);
318 #ifdef FLANN_KDTREE_MEM_OPT
329 void findNeighbors(ResultSet<DistanceType>& result,
const ElementType* vec,
const SearchParams& searchParams, Heap<BranchSt>* heap)
const
331 int maxChecks = searchParams.checks;
332 float epsError = 1+searchParams.eps;
336 getExactNeighbors<true>(result, vec, epsError);
339 getExactNeighbors<false>(result, vec, epsError);
344 getNeighbors<true>(result, vec, maxChecks, epsError, heap);
347 getNeighbors<false>(result, vec, maxChecks, epsError, heap);
364 const SearchParams& params)
const
370 assert(dists.
cols >= knn);
381 Heap<BranchSt>* heap =
new Heap<BranchSt>((
int)
size_);
386 KNNResultSet2<DistanceType> resultSet(knn);
388 for (
int i = 0;
i < (
int)queries.
rows; i++) {
391 size_t n =
std::min(resultSet.size(), knn);
392 resultSet.copy(indices[i], dists[i], n,
params.sorted);
399 std::vector<double> times(queries.
rows);
402 KNNSimpleResultSet<DistanceType> resultSet(knn);
404 for (
int i = 0;
i < (
int)queries.
rows; i++) {
407 size_t n =
std::min(resultSet.size(), knn);
408 resultSet.copy(indices[i], dists[i], n,
params.sorted);
413 std::sort(times.begin(), times.end());
429 std::vector< std::vector<size_t> >& indices,
430 std::vector<std::vector<DistanceType> >& dists,
432 const SearchParams& params)
const
446 Heap<BranchSt>* heap =
new Heap<BranchSt>((
int)
size_);
452 KNNResultSet2<DistanceType> resultSet(knn);
454 for (
int i = 0;
i < (
int)queries.
rows; i++) {
457 size_t n =
std::min(resultSet.size(), knn);
461 resultSet.copy(&indices[i][0], &dists[i][0], n,
params.sorted);
471 KNNSimpleResultSet<DistanceType> resultSet(knn);
473 for (
int i = 0;
i < (
int)queries.
rows; i++) {
476 size_t n =
std::min(resultSet.size(), knn);
480 resultSet.copy(&indices[i][0], &dists[i][0], n,
params.sorted);
505 const SearchParams& params)
const
510 int max_neighbors =
params.max_neighbors;
511 if (max_neighbors<0) max_neighbors = num_neighbors;
512 else max_neighbors =
std::min(max_neighbors,(
int)num_neighbors);
514 Heap<BranchSt>* heap =
new Heap<BranchSt>((
int)
size_);
516 if (max_neighbors==0) {
519 CountRadiusResultSet<DistanceType> resultSet(radius);
521 for (
int i = 0;
i < (
int)queries.
rows; i++) {
524 count += resultSet.size();
531 if (
params.max_neighbors<0 && (num_neighbors>=
this->size())) {
534 RadiusResultSet<DistanceType> resultSet(radius);
536 for (
int i = 0;
i < (
int)queries.
rows; i++) {
539 size_t n = resultSet.size();
541 if (n>num_neighbors)
n = num_neighbors;
542 resultSet.copy(indices[i], dists[i], n,
params.sorted);
546 if (n<dists.
cols) dists[
i][
n] = std::numeric_limits<DistanceType>::infinity();
555 KNNRadiusResultSet<DistanceType> resultSet(radius, max_neighbors);
557 for (
int i = 0;
i < (
int)queries.
rows; i++) {
560 size_t n = resultSet.size();
562 if ((
int)n>max_neighbors)
n = max_neighbors;
563 resultSet.copy(indices[i], dists[i], n,
params.sorted);
567 if (n<dists.
cols) dists[
i][
n] = std::numeric_limits<DistanceType>::infinity();
587 std::vector< std::vector<size_t> >& indices,
588 std::vector<std::vector<DistanceType> >& dists,
590 const SearchParams& params)
const
595 Heap<BranchSt>* heap =
new Heap<BranchSt>((
int)
size_);
598 if (
params.max_neighbors==0) {
601 CountRadiusResultSet<DistanceType> resultSet(radius);
603 for (
int i = 0;
i < (
int)queries.
rows; i++) {
606 count += resultSet.size();
614 if (
params.max_neighbors<0) {
618 RadiusResultSet<DistanceType> resultSet(radius);
620 for (
int i = 0;
i < (
int)queries.
rows; i++) {
623 size_t n = resultSet.size();
628 resultSet.copy(&indices[i][0], &dists[i][0], n,
params.sorted);
638 KNNRadiusResultSet<DistanceType> resultSet(radius,
params.max_neighbors);
640 for (
int i = 0;
i < (
int)queries.
rows; i++) {
643 size_t n = resultSet.size();
649 resultSet.copy(&indices[i][0], &dists[i][0], n,
params.sorted);
670 for (
size_t i = 0;
i <
size_; ++
i) {
681 std::random_device rd;
682 std::mt19937
g(rd());
683 std::shuffle(
ind.begin(),
ind.end(), g);
705 dst->divfeat = src->divfeat;
706 dst->divval = src->divval;
707 if (src->child1==
NULL && src->child2==
NULL) {
734 node->divfeat = *
ind;
741 meanSplit(ind, count, idx, cutfeat, cutval);
743 node->divfeat = cutfeat;
744 node->divval = cutval;
746 node->child2 =
divideTree(ind+idx, count-idx);
767 for (
int j = 0;
j < cnt; ++
j) {
769 for (
size_t k=0; k<
veclen_; ++k) {
774 for (
size_t k=0; k<
veclen_; ++k) {
775 mean_[k] *= div_factor;
779 for (
int j = 0;
j < cnt; ++
j) {
781 for (
size_t k=0; k<
veclen_; ++k) {
788 cutval =
mean_[cutfeat];
791 planeSplit(ind, count, cutfeat, cutval, lim1, lim2);
793 if (lim1>count/2) index = lim1;
794 else if (lim2<count/2) index = lim2;
795 else index =
count/2;
800 if ((lim1==count)||(lim2==0)) index =
count/2;
815 if ((num <
RAND_DIM)||(v[i] > v[topind[num-1]])) {
825 while (j > 0 && v[topind[j]] > v[topind[j-1]]) {
833 return (
int)topind[rnd];
852 while (left<=right &&
points_[
ind[left]][cutfeat]<cutval) ++left;
853 while (left<=right &&
points_[
ind[right]][cutfeat]>=cutval) --right;
854 if (left>right)
break;
860 while (left<=right &&
points_[
ind[left]][cutfeat]<=cutval) ++left;
861 while (left<=right &&
points_[
ind[right]][cutfeat]>cutval) --right;
862 if (left>right)
break;
872 template<
bool with_removed>
878 fprintf(
stderr,
"It doesn't make any sense to use more than one tree for exact search");
890 template<
bool with_removed>
902 searchLevel<with_removed>(
result, vec,
tree_roots_[
i], 0, checkCount, maxCheck, epsError, heap, checked);
906 while ( heap->
popMin(branch) && (checkCount < maxCheck || !
result.full() )) {
907 searchLevel<with_removed>(result, vec, branch.node, branch.mindist, checkCount, maxCheck, epsError, heap, checked);
914 #ifdef FLANN_KDTREE_MEM_OPT
920 template<
bool with_removed>
921 void getNeighbors(ResultSet<DistanceType>& result,
const ElementType* vec,
int maxCheck,
float epsError, Heap<BranchSt>* heap)
const
927 DynamicBitset checked(
size_);
932 searchLevel<with_removed>(result, vec,
tree_roots_[i], 0, checkCount, maxCheck, epsError, heap, checked);
936 while ( heap->popMin(branch) && (checkCount < maxCheck || !
result.full() )) {
937 searchLevel<with_removed>(result, vec, branch.node, branch.mindist, checkCount, maxCheck, epsError, heap, checked);
948 template<
bool with_removed>
950 float epsError, Heap<BranchSt>* heap, DynamicBitset& checked)
const
952 if (result_set.worstDist()<mindist) {
958 if ((node->child1 == NULL)&&(node->child2 == NULL)) {
959 int index = node->divfeat;
964 if ( checked.test(index) || ((checkCount>=maxCheck)&& result_set.full()) )
return;
969 result_set.addPoint(dist,index);
989 if ((new_distsq*epsError < result_set.worstDist())|| !result_set.full()) {
990 heap->insert(
BranchSt(otherChild, new_distsq) );
994 searchLevel<with_removed>(result_set, vec, bestChild, mindist, checkCount, maxCheck, epsError, heap, checked);
1000 template<
bool with_removed>
1004 if ((node->child1 == NULL)&&(node->child2 == NULL)) {
1005 int index = node->divfeat;
1010 result_set.addPoint(dist,index);
1032 searchLevelExact<with_removed>(result_set, vec, bestChild, mindist, epsError);
1034 if (mindist*epsError<=result_set.worstDist()) {
1035 searchLevelExact<with_removed>(result_set, vec, otherChild, new_distsq, epsError);
1043 if ((node->child1==NULL) && (node->child2==NULL)) {
1046 size_t div_feat = 0;
1049 if (span > max_span) {
1059 if (point[div_feat]<leaf_point[div_feat]) {
1060 left->divfeat =
ind;
1061 left->point =
point;
1062 right->divfeat = node->divfeat;
1063 right->point = node->point;
1066 left->divfeat = node->divfeat;
1067 left->point = node->point;
1068 right->divfeat =
ind;
1069 right->point =
point;
1071 node->divfeat = div_feat;
1072 node->divval = (
point[div_feat]+leaf_point[div_feat])/2;
1073 node->child1 = left;
1074 node->child2 = right;
1077 if (
point[node->divfeat]<node->divval) {
1142 #endif //FLANN_KDTREE_INDEX_H_