nanoflann
C++ header-only ANN library
nanoflann.hpp
1 /***********************************************************************
2  * Software License Agreement (BSD License)
3  *
4  * Copyright 2008-2009 Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
5  * Copyright 2008-2009 David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
6  * Copyright 2011-2022 Jose Luis Blanco (joseluisblancoc@gmail.com).
7  * All rights reserved.
8  *
9  * THE BSD LICENSE
10  *
11  * Redistribution and use in source and binary forms, with or without
12  * modification, are permitted provided that the following conditions
13  * are met:
14  *
15  * 1. Redistributions of source code must retain the above copyright
16  * notice, this list of conditions and the following disclaimer.
17  * 2. Redistributions in binary form must reproduce the above copyright
18  * notice, this list of conditions and the following disclaimer in the
19  * documentation and/or other materials provided with the distribution.
20  *
21  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
22  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
23  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
24  * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
25  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
26  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
27  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
28  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
29  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
30  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
31  *************************************************************************/
32 
45 #pragma once
46 
47 #include <algorithm>
48 #include <array>
49 #include <cassert>
50 #include <cmath> // for abs()
51 #include <cstdlib> // for abs()
52 #include <functional>
53 #include <istream>
54 #include <limits> // std::reference_wrapper
55 #include <ostream>
56 #include <stdexcept>
57 #include <unordered_set>
58 #include <vector>
59 
61 #define NANOFLANN_VERSION 0x142
62 
63 // Avoid conflicting declaration of min/max macros in windows headers
64 #if !defined(NOMINMAX) && \
65  (defined(_WIN32) || defined(_WIN32_) || defined(WIN32) || defined(_WIN64))
66 #define NOMINMAX
67 #ifdef max
68 #undef max
69 #undef min
70 #endif
71 #endif
72 
73 namespace nanoflann
74 {
79 template <typename T>
81 {
82  return static_cast<T>(3.14159265358979323846);
83 }
84 
89 template <typename T, typename = int>
90 struct has_resize : std::false_type
91 {
92 };
93 
94 template <typename T>
95 struct has_resize<T, decltype((void)std::declval<T>().resize(1), 0)>
96  : std::true_type
97 {
98 };
99 
100 template <typename T, typename = int>
101 struct has_assign : std::false_type
102 {
103 };
104 
105 template <typename T>
106 struct has_assign<T, decltype((void)std::declval<T>().assign(1, 0), 0)>
107  : std::true_type
108 {
109 };
110 
114 template <typename Container>
115 inline typename std::enable_if<has_resize<Container>::value, void>::type resize(
116  Container& c, const size_t nElements)
117 {
118  c.resize(nElements);
119 }
120 
125 template <typename Container>
126 inline typename std::enable_if<!has_resize<Container>::value, void>::type
127  resize(Container& c, const size_t nElements)
128 {
129  if (nElements != c.size())
130  throw std::logic_error("Try to change the size of a std::array.");
131 }
132 
136 template <typename Container, typename T>
137 inline typename std::enable_if<has_assign<Container>::value, void>::type assign(
138  Container& c, const size_t nElements, const T& value)
139 {
140  c.assign(nElements, value);
141 }
142 
146 template <typename Container, typename T>
147 inline typename std::enable_if<!has_assign<Container>::value, void>::type
148  assign(Container& c, const size_t nElements, const T& value)
149 {
150  for (size_t i = 0; i < nElements; i++) c[i] = value;
151 }
152 
155 template <
156  typename _DistanceType, typename _IndexType = size_t,
157  typename _CountType = size_t>
159 {
160  public:
161  using DistanceType = _DistanceType;
162  using IndexType = _IndexType;
163  using CountType = _CountType;
164 
165  private:
166  IndexType* indices;
167  DistanceType* dists;
168  CountType capacity;
169  CountType count;
170 
171  public:
172  explicit inline KNNResultSet(CountType capacity_)
173  : indices(0), dists(0), capacity(capacity_), count(0)
174  {
175  }
176 
177  inline void init(IndexType* indices_, DistanceType* dists_)
178  {
179  indices = indices_;
180  dists = dists_;
181  count = 0;
182  if (capacity)
183  dists[capacity - 1] = (std::numeric_limits<DistanceType>::max)();
184  }
185 
186  inline CountType size() const { return count; }
187 
188  inline bool full() const { return count == capacity; }
189 
195  inline bool addPoint(DistanceType dist, IndexType index)
196  {
197  CountType i;
198  for (i = count; i > 0; --i)
199  {
200 #ifdef NANOFLANN_FIRST_MATCH // If defined and two points have the same
201  // distance, the one with the lowest-index will be
202  // returned first.
203  if ((dists[i - 1] > dist) ||
204  ((dist == dists[i - 1]) && (indices[i - 1] > index)))
205  {
206 #else
207  if (dists[i - 1] > dist)
208  {
209 #endif
210  if (i < capacity)
211  {
212  dists[i] = dists[i - 1];
213  indices[i] = indices[i - 1];
214  }
215  }
216  else
217  break;
218  }
219  if (i < capacity)
220  {
221  dists[i] = dist;
222  indices[i] = index;
223  }
224  if (count < capacity) count++;
225 
226  // tell caller that the search shall continue
227  return true;
228  }
229 
230  inline DistanceType worstDist() const { return dists[capacity - 1]; }
231 };
232 
235 {
237  template <typename PairType>
238  inline bool operator()(const PairType& p1, const PairType& p2) const
239  {
240  return p1.second < p2.second;
241  }
242 };
243 
247 template <typename _DistanceType, typename _IndexType = size_t>
249 {
250  public:
251  using DistanceType = _DistanceType;
252  using IndexType = _IndexType;
253 
254  public:
255  const DistanceType radius;
256 
257  std::vector<std::pair<IndexType, DistanceType>>& m_indices_dists;
258 
259  explicit inline RadiusResultSet(
260  DistanceType radius_,
261  std::vector<std::pair<IndexType, DistanceType>>& indices_dists)
262  : radius(radius_), m_indices_dists(indices_dists)
263  {
264  init();
265  }
266 
267  inline void init() { clear(); }
268  inline void clear() { m_indices_dists.clear(); }
269 
270  inline size_t size() const { return m_indices_dists.size(); }
271 
272  inline bool full() const { return true; }
273 
279  inline bool addPoint(DistanceType dist, IndexType index)
280  {
281  if (dist < radius)
282  m_indices_dists.push_back(std::make_pair(index, dist));
283  return true;
284  }
285 
286  inline DistanceType worstDist() const { return radius; }
287 
292  std::pair<IndexType, DistanceType> worst_item() const
293  {
294  if (m_indices_dists.empty())
295  throw std::runtime_error(
296  "Cannot invoke RadiusResultSet::worst_item() on "
297  "an empty list of results.");
298  using DistIt = typename std::vector<
299  std::pair<IndexType, DistanceType>>::const_iterator;
300  DistIt it = std::max_element(
301  m_indices_dists.begin(), m_indices_dists.end(), IndexDist_Sorter());
302  return *it;
303  }
304 };
305 
310 template <typename T>
311 void save_value(std::ostream& stream, const T& value)
312 {
313  stream.write(reinterpret_cast<const char*>(&value), sizeof(T));
314 }
315 
316 template <typename T>
317 void save_value(std::ostream& stream, const std::vector<T>& value)
318 {
319  size_t size = value.size();
320  stream.write(reinterpret_cast<const char*>(&size), sizeof(size_t));
321  stream.write(reinterpret_cast<const char*>(value.data()), sizeof(T) * size);
322 }
323 
324 template <typename T>
325 void load_value(std::istream& stream, T& value)
326 {
327  stream.read(reinterpret_cast<char*>(&value), sizeof(T));
328 }
329 
330 template <typename T>
331 void load_value(std::istream& stream, std::vector<T>& value)
332 {
333  size_t size;
334  stream.read(reinterpret_cast<char*>(&size), sizeof(size_t));
335  value.resize(size);
336  stream.read(reinterpret_cast<char*>(value.data()), sizeof(T) * size);
337 }
343 struct Metric
344 {
345 };
346 
357 template <
358  class T, class DataSource, typename _DistanceType = T,
359  typename AccessorType = uint32_t>
361 {
362  using ElementType = T;
363  using DistanceType = _DistanceType;
364 
365  const DataSource& data_source;
366 
367  L1_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
368 
369  inline DistanceType evalMetric(
370  const T* a, const AccessorType b_idx, size_t size,
371  DistanceType worst_dist = -1) const
372  {
373  DistanceType result = DistanceType();
374  const T* last = a + size;
375  const T* lastgroup = last - 3;
376  size_t d = 0;
377 
378  /* Process 4 items with each loop for efficiency. */
379  while (a < lastgroup)
380  {
381  const DistanceType diff0 =
382  std::abs(a[0] - data_source.kdtree_get_pt(b_idx, d++));
383  const DistanceType diff1 =
384  std::abs(a[1] - data_source.kdtree_get_pt(b_idx, d++));
385  const DistanceType diff2 =
386  std::abs(a[2] - data_source.kdtree_get_pt(b_idx, d++));
387  const DistanceType diff3 =
388  std::abs(a[3] - data_source.kdtree_get_pt(b_idx, d++));
389  result += diff0 + diff1 + diff2 + diff3;
390  a += 4;
391  if ((worst_dist > 0) && (result > worst_dist)) { return result; }
392  }
393  /* Process last 0-3 components. Not needed for standard vector lengths.
394  */
395  while (a < last)
396  {
397  result += std::abs(*a++ - data_source.kdtree_get_pt(b_idx, d++));
398  }
399  return result;
400  }
401 
402  template <typename U, typename V>
403  inline DistanceType accum_dist(const U a, const V b, const size_t) const
404  {
405  return std::abs(a - b);
406  }
407 };
408 
419 template <
420  class T, class DataSource, typename _DistanceType = T,
421  typename AccessorType = uint32_t>
423 {
424  using ElementType = T;
425  using DistanceType = _DistanceType;
426 
427  const DataSource& data_source;
428 
429  L2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
430 
431  inline DistanceType evalMetric(
432  const T* a, const AccessorType b_idx, size_t size,
433  DistanceType worst_dist = -1) const
434  {
435  DistanceType result = DistanceType();
436  const T* last = a + size;
437  const T* lastgroup = last - 3;
438  size_t d = 0;
439 
440  /* Process 4 items with each loop for efficiency. */
441  while (a < lastgroup)
442  {
443  const DistanceType diff0 =
444  a[0] - data_source.kdtree_get_pt(b_idx, d++);
445  const DistanceType diff1 =
446  a[1] - data_source.kdtree_get_pt(b_idx, d++);
447  const DistanceType diff2 =
448  a[2] - data_source.kdtree_get_pt(b_idx, d++);
449  const DistanceType diff3 =
450  a[3] - data_source.kdtree_get_pt(b_idx, d++);
451  result +=
452  diff0 * diff0 + diff1 * diff1 + diff2 * diff2 + diff3 * diff3;
453  a += 4;
454  if ((worst_dist > 0) && (result > worst_dist)) { return result; }
455  }
456  /* Process last 0-3 components. Not needed for standard vector lengths.
457  */
458  while (a < last)
459  {
460  const DistanceType diff0 =
461  *a++ - data_source.kdtree_get_pt(b_idx, d++);
462  result += diff0 * diff0;
463  }
464  return result;
465  }
466 
467  template <typename U, typename V>
468  inline DistanceType accum_dist(const U a, const V b, const size_t) const
469  {
470  return (a - b) * (a - b);
471  }
472 };
473 
484 template <
485  class T, class DataSource, typename _DistanceType = T,
486  typename AccessorType = uint32_t>
488 {
489  using ElementType = T;
490  using DistanceType = _DistanceType;
491 
492  const DataSource& data_source;
493 
494  L2_Simple_Adaptor(const DataSource& _data_source)
495  : data_source(_data_source)
496  {
497  }
498 
499  inline DistanceType evalMetric(
500  const T* a, const AccessorType b_idx, size_t size) const
501  {
502  DistanceType result = DistanceType();
503  for (size_t i = 0; i < size; ++i)
504  {
505  const DistanceType diff =
506  a[i] - data_source.kdtree_get_pt(b_idx, i);
507  result += diff * diff;
508  }
509  return result;
510  }
511 
512  template <typename U, typename V>
513  inline DistanceType accum_dist(const U a, const V b, const size_t) const
514  {
515  return (a - b) * (a - b);
516  }
517 };
518 
529 template <
530  class T, class DataSource, typename _DistanceType = T,
531  typename AccessorType = uint32_t>
533 {
534  using ElementType = T;
535  using DistanceType = _DistanceType;
536 
537  const DataSource& data_source;
538 
539  SO2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
540 
541  inline DistanceType evalMetric(
542  const T* a, const AccessorType b_idx, size_t size) const
543  {
544  return accum_dist(
545  a[size - 1], data_source.kdtree_get_pt(b_idx, size - 1), size - 1);
546  }
547 
550  template <typename U, typename V>
551  inline DistanceType accum_dist(const U a, const V b, const size_t) const
552  {
553  DistanceType result = DistanceType();
554  DistanceType PI = pi_const<DistanceType>();
555  result = b - a;
556  if (result > PI)
557  result -= 2 * PI;
558  else if (result < -PI)
559  result += 2 * PI;
560  return result;
561  }
562 };
563 
574 template <
575  class T, class DataSource, typename _DistanceType = T,
576  typename AccessorType = uint32_t>
578 {
579  using ElementType = T;
580  using DistanceType = _DistanceType;
581 
583  distance_L2_Simple;
584 
585  SO3_Adaptor(const DataSource& _data_source)
586  : distance_L2_Simple(_data_source)
587  {
588  }
589 
590  inline DistanceType evalMetric(
591  const T* a, const AccessorType b_idx, size_t size) const
592  {
593  return distance_L2_Simple.evalMetric(a, b_idx, size);
594  }
595 
596  template <typename U, typename V>
597  inline DistanceType accum_dist(const U a, const V b, const size_t idx) const
598  {
599  return distance_L2_Simple.accum_dist(a, b, idx);
600  }
601 };
602 
604 struct metric_L1 : public Metric
605 {
606  template <class T, class DataSource, typename AccessorType = uint32_t>
607  struct traits
608  {
610  };
611 };
613 struct metric_L2 : public Metric
614 {
615  template <class T, class DataSource, typename AccessorType = uint32_t>
616  struct traits
617  {
619  };
620 };
622 struct metric_L2_Simple : public Metric
623 {
624  template <class T, class DataSource, typename AccessorType = uint32_t>
625  struct traits
626  {
628  };
629 };
631 struct metric_SO2 : public Metric
632 {
633  template <class T, class DataSource, typename AccessorType = uint32_t>
634  struct traits
635  {
637  };
638 };
640 struct metric_SO3 : public Metric
641 {
642  template <class T, class DataSource, typename AccessorType = uint32_t>
643  struct traits
644  {
646  };
647 };
648 
654 enum class KDTreeSingleIndexAdaptorFlags
655 {
656  None = 0,
657  SkipInitialBuildIndex = 1
658 };
659 
660 inline std::underlying_type<KDTreeSingleIndexAdaptorFlags>::type operator&(
661  KDTreeSingleIndexAdaptorFlags lhs, KDTreeSingleIndexAdaptorFlags rhs)
662 {
663  using underlying =
664  typename std::underlying_type<KDTreeSingleIndexAdaptorFlags>::type;
665  return static_cast<underlying>(lhs) & static_cast<underlying>(rhs);
666 }
667 
670 {
672  size_t _leaf_max_size = 10, KDTreeSingleIndexAdaptorFlags _flags =
673  KDTreeSingleIndexAdaptorFlags::None)
674  : leaf_max_size(_leaf_max_size), flags(_flags)
675  {
676  }
677 
678  size_t leaf_max_size;
679  KDTreeSingleIndexAdaptorFlags flags;
680 };
681 
684 {
687  SearchParams(int checks_IGNORED_ = 32, float eps_ = 0, bool sorted_ = true)
688  : checks(checks_IGNORED_), eps(eps_), sorted(sorted_)
689  {
690  }
691 
692  int checks;
694  float eps;
695  bool sorted;
697 };
710 template <typename T>
711 inline T* allocate(size_t count = 1)
712 {
713  T* mem = static_cast<T*>(::malloc(sizeof(T) * count));
714  return mem;
715 }
716 
732 const size_t WORDSIZE = 16;
733 const size_t BLOCKSIZE = 8192;
734 
736 {
737  /* We maintain memory alignment to word boundaries by requiring that all
738  allocations be in multiples of the machine wordsize. */
739  /* Size of machine word in bytes. Must be power of 2. */
740  /* Minimum number of bytes requested at a time from the system. Must be
741  * multiple of WORDSIZE. */
742 
743  using Offset = uint32_t;
744  using Size = uint32_t;
745  using Dimension = int32_t;
746 
747  Size remaining; /* Number of bytes left in current block of storage. */
748  void* base; /* Pointer to base of current block of storage. */
749  void* loc; /* Current location in block to next allocate memory. */
750 
751  void internal_init()
752  {
753  remaining = 0;
754  base = nullptr;
755  usedMemory = 0;
756  wastedMemory = 0;
757  }
758 
759  public:
760  Size usedMemory;
761  Size wastedMemory;
762 
766  PooledAllocator() { internal_init(); }
767 
771  ~PooledAllocator() { free_all(); }
772 
774  void free_all()
775  {
776  while (base != nullptr)
777  {
778  void* prev =
779  *(static_cast<void**>(base)); /* Get pointer to prev block. */
780  ::free(base);
781  base = prev;
782  }
783  internal_init();
784  }
785 
790  void* malloc(const size_t req_size)
791  {
792  /* Round size up to a multiple of wordsize. The following expression
793  only works for WORDSIZE that is a power of 2, by masking last bits
794  of incremented size to zero.
795  */
796  const Size size = (req_size + (WORDSIZE - 1)) & ~(WORDSIZE - 1);
797 
798  /* Check whether a new block must be allocated. Note that the first
799  word of a block is reserved for a pointer to the previous block.
800  */
801  if (size > remaining)
802  {
803  wastedMemory += remaining;
804 
805  /* Allocate new storage. */
806  const Size blocksize =
807  (size + sizeof(void*) + (WORDSIZE - 1) > BLOCKSIZE)
808  ? size + sizeof(void*) + (WORDSIZE - 1)
809  : BLOCKSIZE;
810 
811  // use the standard C malloc to allocate memory
812  void* m = ::malloc(blocksize);
813  if (!m)
814  {
815  fprintf(stderr, "Failed to allocate memory.\n");
816  throw std::bad_alloc();
817  }
818 
819  /* Fill first word of new block with pointer to previous block. */
820  static_cast<void**>(m)[0] = base;
821  base = m;
822 
823  Size shift = 0;
824  // int size_t = (WORDSIZE - ( (((size_t)m) + sizeof(void*)) &
825  // (WORDSIZE-1))) & (WORDSIZE-1);
826 
827  remaining = blocksize - sizeof(void*) - shift;
828  loc = (static_cast<char*>(m) + sizeof(void*) + shift);
829  }
830  void* rloc = loc;
831  loc = static_cast<char*>(loc) + size;
832  remaining -= size;
833 
834  usedMemory += size;
835 
836  return rloc;
837  }
838 
846  template <typename T>
847  T* allocate(const size_t count = 1)
848  {
849  T* mem = static_cast<T*>(this->malloc(sizeof(T) * count));
850  return mem;
851  }
852 };
861 template <int32_t DIM, typename T>
863 {
864  using container_t = std::array<T, DIM>;
865 };
867 template <typename T>
869 {
870  using container_t = std::vector<T>;
871 };
872 
888 template <
889  class Derived, typename Distance, class DatasetAdaptor, int32_t DIM = -1,
890  typename AccessorType = uint32_t>
892 {
893  public:
896  void freeIndex(Derived& obj)
897  {
898  obj.pool.free_all();
899  obj.root_node = nullptr;
900  obj.m_size_at_index_build = 0;
901  }
902 
903  using ElementType = typename Distance::ElementType;
904  using DistanceType = typename Distance::DistanceType;
905 
909  std::vector<AccessorType> vAcc;
910 
911  using Offset = typename decltype(vAcc)::size_type;
912  using Size = typename decltype(vAcc)::size_type;
913  using Dimension = int32_t;
914 
915  /*--------------------- Internal Data Structures
916  * --------------------------*/
917  struct Node
918  {
921  union
922  {
923  struct leaf
924  {
925  Offset left, right;
926  } lr;
927  struct nonleaf
928  {
929  Dimension divfeat;
930  DistanceType divlow,
932  } sub;
933  } node_type;
935  Node *child1, *child2;
936  };
937 
938  using NodePtr = Node*;
939 
940  struct Interval
941  {
942  ElementType low, high;
943  };
944 
945  NodePtr root_node;
946 
947  Size m_leaf_max_size;
948 
949  Size m_size;
952  Dimension dim;
953 
956  using BoundingBox =
958 
962  typename array_or_vector_selector<DIM, DistanceType>::container_t;
963 
967 
976 
978  Size size(const Derived& obj) const { return obj.m_size; }
979 
981  Size veclen(const Derived& obj) { return DIM > 0 ? DIM : obj.dim; }
982 
984  inline ElementType dataset_get(
985  const Derived& obj, AccessorType element, Dimension component) const
986  {
987  return obj.dataset.kdtree_get_pt(element, component);
988  }
989 
994  Size usedMemory(Derived& obj)
995  {
996  return obj.pool.usedMemory + obj.pool.wastedMemory +
997  obj.dataset.kdtree_get_point_count() *
998  sizeof(AccessorType); // pool memory and vind array memory
999  }
1000 
1001  void computeMinMax(
1002  const Derived& obj, Offset ind, Size count, Dimension element,
1003  ElementType& min_elem, ElementType& max_elem)
1004  {
1005  min_elem = dataset_get(obj, vAcc[ind], element);
1006  max_elem = min_elem;
1007  for (Offset i = 1; i < count; ++i)
1008  {
1009  ElementType val = dataset_get(obj, vAcc[ind + i], element);
1010  if (val < min_elem) min_elem = val;
1011  if (val > max_elem) max_elem = val;
1012  }
1013  }
1014 
1023  Derived& obj, const Offset left, const Offset right, BoundingBox& bbox)
1024  {
1025  NodePtr node = obj.pool.template allocate<Node>(); // allocate memory
1026 
1027  /* If too few exemplars remain, then make this a leaf node. */
1028  if ((right - left) <= static_cast<Offset>(obj.m_leaf_max_size))
1029  {
1030  node->child1 = node->child2 = nullptr; /* Mark as leaf node. */
1031  node->node_type.lr.left = left;
1032  node->node_type.lr.right = right;
1033 
1034  // compute bounding-box of leaf points
1035  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1036  {
1037  bbox[i].low = dataset_get(obj, obj.vAcc[left], i);
1038  bbox[i].high = dataset_get(obj, obj.vAcc[left], i);
1039  }
1040  for (Offset k = left + 1; k < right; ++k)
1041  {
1042  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1043  {
1044  if (bbox[i].low > dataset_get(obj, obj.vAcc[k], i))
1045  bbox[i].low = dataset_get(obj, obj.vAcc[k], i);
1046  if (bbox[i].high < dataset_get(obj, obj.vAcc[k], i))
1047  bbox[i].high = dataset_get(obj, obj.vAcc[k], i);
1048  }
1049  }
1050  }
1051  else
1052  {
1053  Offset idx;
1054  Dimension cutfeat;
1055  DistanceType cutval;
1056  middleSplit_(obj, left, right - left, idx, cutfeat, cutval, bbox);
1057 
1058  node->node_type.sub.divfeat = cutfeat;
1059 
1060  BoundingBox left_bbox(bbox);
1061  left_bbox[cutfeat].high = cutval;
1062  node->child1 = divideTree(obj, left, left + idx, left_bbox);
1063 
1064  BoundingBox right_bbox(bbox);
1065  right_bbox[cutfeat].low = cutval;
1066  node->child2 = divideTree(obj, left + idx, right, right_bbox);
1067 
1068  node->node_type.sub.divlow = left_bbox[cutfeat].high;
1069  node->node_type.sub.divhigh = right_bbox[cutfeat].low;
1070 
1071  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1072  {
1073  bbox[i].low = std::min(left_bbox[i].low, right_bbox[i].low);
1074  bbox[i].high = std::max(left_bbox[i].high, right_bbox[i].high);
1075  }
1076  }
1077 
1078  return node;
1079  }
1080 
1081  void middleSplit_(
1082  Derived& obj, Offset ind, Size count, Offset& index, Dimension& cutfeat,
1083  DistanceType& cutval, const BoundingBox& bbox)
1084  {
1085  const auto EPS = static_cast<DistanceType>(0.00001);
1086  ElementType max_span = bbox[0].high - bbox[0].low;
1087  for (Dimension i = 1; i < (DIM > 0 ? DIM : obj.dim); ++i)
1088  {
1089  ElementType span = bbox[i].high - bbox[i].low;
1090  if (span > max_span) { max_span = span; }
1091  }
1092  ElementType max_spread = -1;
1093  cutfeat = 0;
1094  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1095  {
1096  ElementType span = bbox[i].high - bbox[i].low;
1097  if (span > (1 - EPS) * max_span)
1098  {
1099  ElementType min_elem, max_elem;
1100  computeMinMax(obj, ind, count, i, min_elem, max_elem);
1101  ElementType spread = max_elem - min_elem;
1102  if (spread > max_spread)
1103  {
1104  cutfeat = i;
1105  max_spread = spread;
1106  }
1107  }
1108  }
1109  // split in the middle
1110  DistanceType split_val = (bbox[cutfeat].low + bbox[cutfeat].high) / 2;
1111  ElementType min_elem, max_elem;
1112  computeMinMax(obj, ind, count, cutfeat, min_elem, max_elem);
1113 
1114  if (split_val < min_elem)
1115  cutval = min_elem;
1116  else if (split_val > max_elem)
1117  cutval = max_elem;
1118  else
1119  cutval = split_val;
1120 
1121  Offset lim1, lim2;
1122  planeSplit(obj, ind, count, cutfeat, cutval, lim1, lim2);
1123 
1124  if (lim1 > count / 2)
1125  index = lim1;
1126  else if (lim2 < count / 2)
1127  index = lim2;
1128  else
1129  index = count / 2;
1130  }
1131 
1142  Derived& obj, Offset ind, const Size count, Dimension cutfeat,
1143  DistanceType& cutval, Offset& lim1, Offset& lim2)
1144  {
1145  /* Move vector indices for left subtree to front of list. */
1146  Offset left = 0;
1147  Offset right = count - 1;
1148  for (;;)
1149  {
1150  while (left <= right &&
1151  dataset_get(obj, vAcc[ind + left], cutfeat) < cutval)
1152  ++left;
1153  while (right && left <= right &&
1154  dataset_get(obj, vAcc[ind + right], cutfeat) >= cutval)
1155  --right;
1156  if (left > right || !right)
1157  break; // "!right" was added to support unsigned Index types
1158  std::swap(vAcc[ind + left], vAcc[ind + right]);
1159  ++left;
1160  --right;
1161  }
1162  /* If either list is empty, it means that all remaining features
1163  * are identical. Split in the middle to maintain a balanced tree.
1164  */
1165  lim1 = left;
1166  right = count - 1;
1167  for (;;)
1168  {
1169  while (left <= right &&
1170  dataset_get(obj, vAcc[ind + left], cutfeat) <= cutval)
1171  ++left;
1172  while (right && left <= right &&
1173  dataset_get(obj, vAcc[ind + right], cutfeat) > cutval)
1174  --right;
1175  if (left > right || !right)
1176  break; // "!right" was added to support unsigned Index types
1177  std::swap(vAcc[ind + left], vAcc[ind + right]);
1178  ++left;
1179  --right;
1180  }
1181  lim2 = left;
1182  }
1183 
1184  DistanceType computeInitialDistances(
1185  const Derived& obj, const ElementType* vec,
1186  distance_vector_t& dists) const
1187  {
1188  assert(vec);
1189  DistanceType distsq = DistanceType();
1190 
1191  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1192  {
1193  if (vec[i] < obj.root_bbox[i].low)
1194  {
1195  dists[i] =
1196  obj.distance.accum_dist(vec[i], obj.root_bbox[i].low, i);
1197  distsq += dists[i];
1198  }
1199  if (vec[i] > obj.root_bbox[i].high)
1200  {
1201  dists[i] =
1202  obj.distance.accum_dist(vec[i], obj.root_bbox[i].high, i);
1203  distsq += dists[i];
1204  }
1205  }
1206  return distsq;
1207  }
1208 
1209  void save_tree(Derived& obj, std::ostream& stream, NodePtr tree)
1210  {
1211  save_value(stream, *tree);
1212  if (tree->child1 != nullptr) { save_tree(obj, stream, tree->child1); }
1213  if (tree->child2 != nullptr) { save_tree(obj, stream, tree->child2); }
1214  }
1215 
1216  void load_tree(Derived& obj, std::istream& stream, NodePtr& tree)
1217  {
1218  tree = obj.pool.template allocate<Node>();
1219  load_value(stream, *tree);
1220  if (tree->child1 != nullptr) { load_tree(obj, stream, tree->child1); }
1221  if (tree->child2 != nullptr) { load_tree(obj, stream, tree->child2); }
1222  }
1223 
1229  void saveIndex_(Derived& obj, std::ostream& stream)
1230  {
1231  save_value(stream, obj.m_size);
1232  save_value(stream, obj.dim);
1233  save_value(stream, obj.root_bbox);
1234  save_value(stream, obj.m_leaf_max_size);
1235  save_value(stream, obj.vAcc);
1236  save_tree(obj, stream, obj.root_node);
1237  }
1238 
1244  void loadIndex_(Derived& obj, std::istream& stream)
1245  {
1246  load_value(stream, obj.m_size);
1247  load_value(stream, obj.dim);
1248  load_value(stream, obj.root_bbox);
1249  load_value(stream, obj.m_leaf_max_size);
1250  load_value(stream, obj.vAcc);
1251  load_tree(obj, stream, obj.root_node);
1252  }
1253 };
1254 
1296 template <
1297  typename Distance, class DatasetAdaptor, int32_t DIM = -1,
1298  typename AccessorType = uint32_t>
1300  : public KDTreeBaseClass<
1301  KDTreeSingleIndexAdaptor<Distance, DatasetAdaptor, DIM, AccessorType>,
1302  Distance, DatasetAdaptor, DIM, AccessorType>
1303 {
1304  public:
1308  Distance, DatasetAdaptor, DIM, AccessorType>&) = delete;
1309 
1313  const DatasetAdaptor& dataset;
1314 
1315  const KDTreeSingleIndexAdaptorParams index_params;
1316 
1317  Distance distance;
1318 
1319  using BaseClassRef = typename nanoflann::KDTreeBaseClass<
1321  Distance, DatasetAdaptor, DIM, AccessorType>,
1322  Distance, DatasetAdaptor, DIM, AccessorType>;
1323 
1324  using Offset = typename BaseClassRef::Offset;
1325  using Size = typename BaseClassRef::Size;
1326  using Dimension = typename BaseClassRef::Dimension;
1327 
1328  using ElementType = typename BaseClassRef::ElementType;
1329  using DistanceType = typename BaseClassRef::DistanceType;
1330 
1331  using Node = typename BaseClassRef::Node;
1332  using NodePtr = Node*;
1333 
1334  using Interval = typename BaseClassRef::Interval;
1335 
1338  using BoundingBox = typename BaseClassRef::BoundingBox;
1339 
1342  using distance_vector_t = typename BaseClassRef::distance_vector_t;
1343 
1364  template <class... Args>
1366  const Dimension dimensionality, const DatasetAdaptor& inputData,
1367  const KDTreeSingleIndexAdaptorParams& params, Args&&... args)
1368  : dataset(inputData),
1369  index_params(params),
1370  distance(inputData, std::forward<Args>(args)...)
1371  {
1372  init(dimensionality, params);
1373  }
1374 
1375  explicit KDTreeSingleIndexAdaptor(
1376  const Dimension dimensionality, const DatasetAdaptor& inputData,
1377  const KDTreeSingleIndexAdaptorParams& params = {})
1378  : dataset(inputData), index_params(params), distance(inputData)
1379  {
1380  init(dimensionality, params);
1381  }
1382 
1383  private:
1384  void init(
1385  const Dimension dimensionality,
1386  const KDTreeSingleIndexAdaptorParams& params)
1387  {
1388  BaseClassRef::root_node = nullptr;
1389  BaseClassRef::m_size = dataset.kdtree_get_point_count();
1390  BaseClassRef::m_size_at_index_build = BaseClassRef::m_size;
1391  BaseClassRef::dim = dimensionality;
1392  if (DIM > 0) BaseClassRef::dim = DIM;
1393  BaseClassRef::m_leaf_max_size = params.leaf_max_size;
1394 
1395  if (!(params.flags &
1396  KDTreeSingleIndexAdaptorFlags::SkipInitialBuildIndex))
1397  {
1398  buildIndex();
1399  }
1400  }
1401 
1402  public:
1406  void buildIndex()
1407  {
1408  BaseClassRef::m_size = dataset.kdtree_get_point_count();
1409  BaseClassRef::m_size_at_index_build = BaseClassRef::m_size;
1410  init_vind();
1411  this->freeIndex(*this);
1412  BaseClassRef::m_size_at_index_build = BaseClassRef::m_size;
1413  if (BaseClassRef::m_size == 0) return;
1414  computeBoundingBox(BaseClassRef::root_bbox);
1415  BaseClassRef::root_node = this->divideTree(
1416  *this, 0, BaseClassRef::m_size,
1417  BaseClassRef::root_bbox); // construct the tree
1418  }
1419 
1436  template <typename RESULTSET>
1438  RESULTSET& result, const ElementType* vec,
1439  const SearchParams& searchParams) const
1440  {
1441  assert(vec);
1442  if (this->size(*this) == 0) return false;
1443  if (!BaseClassRef::root_node)
1444  throw std::runtime_error(
1445  "[nanoflann] findNeighbors() called before building the "
1446  "index.");
1447  float epsError = 1 + searchParams.eps;
1448 
1450  dists; // fixed or variable-sized container (depending on DIM)
1451  auto zero = static_cast<decltype(result.worstDist())>(0);
1452  assign(
1453  dists, (DIM > 0 ? DIM : BaseClassRef::dim),
1454  zero); // Fill it with zeros.
1455  DistanceType distsq = this->computeInitialDistances(*this, vec, dists);
1456  searchLevel(
1457  result, vec, BaseClassRef::root_node, distsq, dists,
1458  epsError); // "count_leaf" parameter removed since was neither
1459  // used nor returned to the user.
1460  return result.full();
1461  }
1462 
1473  const ElementType* query_point, const Size num_closest,
1474  AccessorType* out_indices, DistanceType* out_distances_sq,
1475  const int /* nChecks_IGNORED */ = 10) const
1476  {
1478  num_closest);
1479  resultSet.init(out_indices, out_distances_sq);
1480  this->findNeighbors(resultSet, query_point, nanoflann::SearchParams());
1481  return resultSet.size();
1482  }
1483 
1501  const ElementType* query_point, const DistanceType& radius,
1502  std::vector<std::pair<AccessorType, DistanceType>>& IndicesDists,
1503  const SearchParams& searchParams) const
1504  {
1506  radius, IndicesDists);
1507  const Size nFound =
1508  radiusSearchCustomCallback(query_point, resultSet, searchParams);
1509  if (searchParams.sorted)
1510  std::sort(
1511  IndicesDists.begin(), IndicesDists.end(), IndexDist_Sorter());
1512  return nFound;
1513  }
1514 
1520  template <class SEARCH_CALLBACK>
1522  const ElementType* query_point, SEARCH_CALLBACK& resultSet,
1523  const SearchParams& searchParams = SearchParams()) const
1524  {
1525  this->findNeighbors(resultSet, query_point, searchParams);
1526  return resultSet.size();
1527  }
1528 
1531  public:
1534  void init_vind()
1535  {
1536  // Create a permutable array of indices to the input vectors.
1537  BaseClassRef::m_size = dataset.kdtree_get_point_count();
1538  if (BaseClassRef::vAcc.size() != BaseClassRef::m_size)
1539  BaseClassRef::vAcc.resize(BaseClassRef::m_size);
1540  for (Size i = 0; i < BaseClassRef::m_size; i++)
1541  BaseClassRef::vAcc[i] = i;
1542  }
1543 
1544  void computeBoundingBox(BoundingBox& bbox)
1545  {
1546  resize(bbox, (DIM > 0 ? DIM : BaseClassRef::dim));
1547  if (dataset.kdtree_get_bbox(bbox))
1548  {
1549  // Done! It was implemented in derived class
1550  }
1551  else
1552  {
1553  const Size N = dataset.kdtree_get_point_count();
1554  if (!N)
1555  throw std::runtime_error(
1556  "[nanoflann] computeBoundingBox() called but "
1557  "no data points found.");
1558  for (Dimension i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim); ++i)
1559  {
1560  bbox[i].low = bbox[i].high =
1561  this->dataset_get(*this, BaseClassRef::vAcc[0], i);
1562  }
1563  for (Offset k = 1; k < N; ++k)
1564  {
1565  for (Dimension i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim);
1566  ++i)
1567  {
1568  if (this->dataset_get(*this, BaseClassRef::vAcc[k], i) <
1569  bbox[i].low)
1570  bbox[i].low =
1571  this->dataset_get(*this, BaseClassRef::vAcc[k], i);
1572  if (this->dataset_get(*this, BaseClassRef::vAcc[k], i) >
1573  bbox[i].high)
1574  bbox[i].high =
1575  this->dataset_get(*this, BaseClassRef::vAcc[k], i);
1576  }
1577  }
1578  }
1579  }
1580 
1587  template <class RESULTSET>
1589  RESULTSET& result_set, const ElementType* vec, const NodePtr node,
1590  DistanceType mindistsq, distance_vector_t& dists,
1591  const float epsError) const
1592  {
1593  /* If this is a leaf node, then do check and return. */
1594  if ((node->child1 == nullptr) && (node->child2 == nullptr))
1595  {
1596  // count_leaf += (node->lr.right-node->lr.left); // Removed since
1597  // was neither used nor returned to the user.
1598  DistanceType worst_dist = result_set.worstDist();
1599  for (Offset i = node->node_type.lr.left;
1600  i < node->node_type.lr.right; ++i)
1601  {
1602  const AccessorType accessor =
1603  BaseClassRef::vAcc[i]; // reorder... : i;
1604  DistanceType dist = distance.evalMetric(
1605  vec, accessor, (DIM > 0 ? DIM : BaseClassRef::dim));
1606  if (dist < worst_dist)
1607  {
1608  if (!result_set.addPoint(dist, BaseClassRef::vAcc[i]))
1609  {
1610  // the resultset doesn't want to receive any more
1611  // points, we're done searching!
1612  return false;
1613  }
1614  }
1615  }
1616  return true;
1617  }
1618 
1619  /* Which child branch should be taken first? */
1620  Dimension idx = node->node_type.sub.divfeat;
1621  ElementType val = vec[idx];
1622  DistanceType diff1 = val - node->node_type.sub.divlow;
1623  DistanceType diff2 = val - node->node_type.sub.divhigh;
1624 
1625  NodePtr bestChild;
1626  NodePtr otherChild;
1627  DistanceType cut_dist;
1628  if ((diff1 + diff2) < 0)
1629  {
1630  bestChild = node->child1;
1631  otherChild = node->child2;
1632  cut_dist =
1633  distance.accum_dist(val, node->node_type.sub.divhigh, idx);
1634  }
1635  else
1636  {
1637  bestChild = node->child2;
1638  otherChild = node->child1;
1639  cut_dist =
1640  distance.accum_dist(val, node->node_type.sub.divlow, idx);
1641  }
1642 
1643  /* Call recursively to search next level down. */
1644  if (!searchLevel(
1645  result_set, vec, bestChild, mindistsq, dists, epsError))
1646  {
1647  // the resultset doesn't want to receive any more points, we're done
1648  // searching!
1649  return false;
1650  }
1651 
1652  DistanceType dst = dists[idx];
1653  mindistsq = mindistsq + cut_dist - dst;
1654  dists[idx] = cut_dist;
1655  if (mindistsq * epsError <= result_set.worstDist())
1656  {
1657  if (!searchLevel(
1658  result_set, vec, otherChild, mindistsq, dists, epsError))
1659  {
1660  // the resultset doesn't want to receive any more points, we're
1661  // done searching!
1662  return false;
1663  }
1664  }
1665  dists[idx] = dst;
1666  return true;
1667  }
1668 
1669  public:
1675  void saveIndex(std::ostream& stream) { this->saveIndex_(*this, stream); }
1676 
1682  void loadIndex(std::istream& stream) { this->loadIndex_(*this, stream); }
1683 
1684 }; // class KDTree
1685 
1722 template <
1723  typename Distance, class DatasetAdaptor, int32_t DIM = -1,
1724  typename AccessorType = uint32_t>
1726  : public KDTreeBaseClass<
1727  KDTreeSingleIndexDynamicAdaptor_<
1728  Distance, DatasetAdaptor, DIM, AccessorType>,
1729  Distance, DatasetAdaptor, DIM, AccessorType>
1730 {
1731  public:
1735  const DatasetAdaptor& dataset;
1736 
1737  KDTreeSingleIndexAdaptorParams index_params;
1738 
1739  std::vector<int>& treeIndex;
1740 
1741  Distance distance;
1742 
1743  using BaseClassRef = typename nanoflann::KDTreeBaseClass<
1745  Distance, DatasetAdaptor, DIM, AccessorType>,
1746  Distance, DatasetAdaptor, DIM, AccessorType>;
1747 
1748  using ElementType = typename BaseClassRef::ElementType;
1749  using DistanceType = typename BaseClassRef::DistanceType;
1750 
1751  using Offset = typename BaseClassRef::Offset;
1752  using Size = typename BaseClassRef::Size;
1753  using Dimension = typename BaseClassRef::Dimension;
1754 
1755  using Node = typename BaseClassRef::Node;
1756  using NodePtr = Node*;
1757 
1758  using Interval = typename BaseClassRef::Interval;
1761  using BoundingBox = typename BaseClassRef::BoundingBox;
1762 
1765  using distance_vector_t = typename BaseClassRef::distance_vector_t;
1766 
1783  const Dimension dimensionality, const DatasetAdaptor& inputData,
1784  std::vector<int>& treeIndex_,
1785  const KDTreeSingleIndexAdaptorParams& params =
1787  : dataset(inputData),
1788  index_params(params),
1789  treeIndex(treeIndex_),
1790  distance(inputData)
1791  {
1792  BaseClassRef::root_node = nullptr;
1793  BaseClassRef::m_size = 0;
1794  BaseClassRef::m_size_at_index_build = 0;
1795  BaseClassRef::dim = dimensionality;
1796  if (DIM > 0) BaseClassRef::dim = DIM;
1797  BaseClassRef::m_leaf_max_size = params.leaf_max_size;
1798  }
1799 
1802  const KDTreeSingleIndexDynamicAdaptor_& rhs) = default;
1803 
1807  {
1809  std::swap(BaseClassRef::vAcc, tmp.BaseClassRef::vAcc);
1810  std::swap(
1811  BaseClassRef::m_leaf_max_size, tmp.BaseClassRef::m_leaf_max_size);
1812  std::swap(index_params, tmp.index_params);
1813  std::swap(treeIndex, tmp.treeIndex);
1814  std::swap(BaseClassRef::m_size, tmp.BaseClassRef::m_size);
1815  std::swap(
1816  BaseClassRef::m_size_at_index_build,
1817  tmp.BaseClassRef::m_size_at_index_build);
1818  std::swap(BaseClassRef::root_node, tmp.BaseClassRef::root_node);
1819  std::swap(BaseClassRef::root_bbox, tmp.BaseClassRef::root_bbox);
1820  std::swap(BaseClassRef::pool, tmp.BaseClassRef::pool);
1821  return *this;
1822  }
1823 
1827  void buildIndex()
1828  {
1829  BaseClassRef::m_size = BaseClassRef::vAcc.size();
1830  this->freeIndex(*this);
1831  BaseClassRef::m_size_at_index_build = BaseClassRef::m_size;
1832  if (BaseClassRef::m_size == 0) return;
1833  computeBoundingBox(BaseClassRef::root_bbox);
1834  BaseClassRef::root_node = this->divideTree(
1835  *this, 0, BaseClassRef::m_size,
1836  BaseClassRef::root_bbox); // construct the tree
1837  }
1838 
1855  template <typename RESULTSET>
1857  RESULTSET& result, const ElementType* vec,
1858  const SearchParams& searchParams) const
1859  {
1860  assert(vec);
1861  if (this->size(*this) == 0) return false;
1862  if (!BaseClassRef::root_node) return false;
1863  float epsError = 1 + searchParams.eps;
1864 
1865  // fixed or variable-sized container (depending on DIM)
1866  distance_vector_t dists;
1867  // Fill it with zeros.
1868  assign(
1869  dists, (DIM > 0 ? DIM : BaseClassRef::dim),
1870  static_cast<typename distance_vector_t::value_type>(0));
1871  DistanceType distsq = this->computeInitialDistances(*this, vec, dists);
1872  searchLevel(
1873  result, vec, BaseClassRef::root_node, distsq, dists,
1874  epsError); // "count_leaf" parameter removed since was neither
1875  // used nor returned to the user.
1876  return result.full();
1877  }
1878 
1889  const ElementType* query_point, const Size num_closest,
1890  AccessorType* out_indices, DistanceType* out_distances_sq,
1891  const int /* nChecks_IGNORED */ = 10) const
1892  {
1894  num_closest);
1895  resultSet.init(out_indices, out_distances_sq);
1896  this->findNeighbors(resultSet, query_point, nanoflann::SearchParams());
1897  return resultSet.size();
1898  }
1899 
1917  const ElementType* query_point, const DistanceType& radius,
1918  std::vector<std::pair<AccessorType, DistanceType>>& IndicesDists,
1919  const SearchParams& searchParams) const
1920  {
1922  radius, IndicesDists);
1923  const size_t nFound =
1924  radiusSearchCustomCallback(query_point, resultSet, searchParams);
1925  if (searchParams.sorted)
1926  std::sort(
1927  IndicesDists.begin(), IndicesDists.end(), IndexDist_Sorter());
1928  return nFound;
1929  }
1930 
1936  template <class SEARCH_CALLBACK>
1938  const ElementType* query_point, SEARCH_CALLBACK& resultSet,
1939  const SearchParams& searchParams = SearchParams()) const
1940  {
1941  this->findNeighbors(resultSet, query_point, searchParams);
1942  return resultSet.size();
1943  }
1944 
1947  public:
1948  void computeBoundingBox(BoundingBox& bbox)
1949  {
1950  resize(bbox, (DIM > 0 ? DIM : BaseClassRef::dim));
1951 
1952  if (dataset.kdtree_get_bbox(bbox))
1953  {
1954  // Done! It was implemented in derived class
1955  }
1956  else
1957  {
1958  const Size N = BaseClassRef::m_size;
1959  if (!N)
1960  throw std::runtime_error(
1961  "[nanoflann] computeBoundingBox() called but "
1962  "no data points found.");
1963  for (Dimension i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim); ++i)
1964  {
1965  bbox[i].low = bbox[i].high =
1966  this->dataset_get(*this, BaseClassRef::vAcc[0], i);
1967  }
1968  for (Offset k = 1; k < N; ++k)
1969  {
1970  for (Dimension i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim);
1971  ++i)
1972  {
1973  if (this->dataset_get(*this, BaseClassRef::vAcc[k], i) <
1974  bbox[i].low)
1975  bbox[i].low =
1976  this->dataset_get(*this, BaseClassRef::vAcc[k], i);
1977  if (this->dataset_get(*this, BaseClassRef::vAcc[k], i) >
1978  bbox[i].high)
1979  bbox[i].high =
1980  this->dataset_get(*this, BaseClassRef::vAcc[k], i);
1981  }
1982  }
1983  }
1984  }
1985 
1990  template <class RESULTSET>
1992  RESULTSET& result_set, const ElementType* vec, const NodePtr node,
1993  DistanceType mindistsq, distance_vector_t& dists,
1994  const float epsError) const
1995  {
1996  /* If this is a leaf node, then do check and return. */
1997  if ((node->child1 == nullptr) && (node->child2 == nullptr))
1998  {
1999  // count_leaf += (node->lr.right-node->lr.left); // Removed since
2000  // was neither used nor returned to the user.
2001  DistanceType worst_dist = result_set.worstDist();
2002  for (Offset i = node->node_type.lr.left;
2003  i < node->node_type.lr.right; ++i)
2004  {
2005  const AccessorType index =
2006  BaseClassRef::vAcc[i]; // reorder... : i;
2007  if (treeIndex[index] == -1) continue;
2008  DistanceType dist = distance.evalMetric(
2009  vec, index, (DIM > 0 ? DIM : BaseClassRef::dim));
2010  if (dist < worst_dist)
2011  {
2012  if (!result_set.addPoint(
2013  static_cast<typename RESULTSET::DistanceType>(dist),
2014  static_cast<typename RESULTSET::IndexType>(
2015  BaseClassRef::vAcc[i])))
2016  {
2017  // the resultset doesn't want to receive any more
2018  // points, we're done searching!
2019  return; // false;
2020  }
2021  }
2022  }
2023  return;
2024  }
2025 
2026  /* Which child branch should be taken first? */
2027  Dimension idx = node->node_type.sub.divfeat;
2028  ElementType val = vec[idx];
2029  DistanceType diff1 = val - node->node_type.sub.divlow;
2030  DistanceType diff2 = val - node->node_type.sub.divhigh;
2031 
2032  NodePtr bestChild;
2033  NodePtr otherChild;
2034  DistanceType cut_dist;
2035  if ((diff1 + diff2) < 0)
2036  {
2037  bestChild = node->child1;
2038  otherChild = node->child2;
2039  cut_dist =
2040  distance.accum_dist(val, node->node_type.sub.divhigh, idx);
2041  }
2042  else
2043  {
2044  bestChild = node->child2;
2045  otherChild = node->child1;
2046  cut_dist =
2047  distance.accum_dist(val, node->node_type.sub.divlow, idx);
2048  }
2049 
2050  /* Call recursively to search next level down. */
2051  searchLevel(result_set, vec, bestChild, mindistsq, dists, epsError);
2052 
2053  DistanceType dst = dists[idx];
2054  mindistsq = mindistsq + cut_dist - dst;
2055  dists[idx] = cut_dist;
2056  if (mindistsq * epsError <= result_set.worstDist())
2057  {
2058  searchLevel(
2059  result_set, vec, otherChild, mindistsq, dists, epsError);
2060  }
2061  dists[idx] = dst;
2062  }
2063 
2064  public:
2070  void saveIndex(std::ostream& stream) { this->saveIndex_(*this, stream); }
2071 
2077  void loadIndex(std::istream& stream) { this->loadIndex_(*this, stream); }
2078 };
2079 
2094 template <
2095  typename Distance, class DatasetAdaptor, int32_t DIM = -1,
2096  typename AccessorType = uint32_t>
2098 {
2099  public:
2100  using ElementType = typename Distance::ElementType;
2101  using DistanceType = typename Distance::DistanceType;
2102 
2103  using Offset = typename KDTreeSingleIndexDynamicAdaptor_<
2104  Distance, DatasetAdaptor, DIM>::Offset;
2105  using Size = typename KDTreeSingleIndexDynamicAdaptor_<
2106  Distance, DatasetAdaptor, DIM>::Size;
2107  using Dimension = typename KDTreeSingleIndexDynamicAdaptor_<
2108  Distance, DatasetAdaptor, DIM>::Dimension;
2109 
2110  protected:
2111  Size m_leaf_max_size;
2112  Size treeCount;
2113  Size pointCount;
2114 
2118  const DatasetAdaptor& dataset;
2119 
2120  std::vector<int>
2124  std::unordered_set<int> removedPoints;
2125 
2126  KDTreeSingleIndexAdaptorParams index_params;
2127 
2128  Dimension dim;
2129 
2131  Distance, DatasetAdaptor, DIM, AccessorType>;
2132  std::vector<index_container_t> index;
2133 
2134  public:
2137  const std::vector<index_container_t>& getAllIndices() const
2138  {
2139  return index;
2140  }
2141 
2142  private:
2144  int First0Bit(AccessorType num)
2145  {
2146  int pos = 0;
2147  while (num & 1)
2148  {
2149  num = num >> 1;
2150  pos++;
2151  }
2152  return pos;
2153  }
2154 
2156  void init()
2157  {
2158  using my_kd_tree_t = KDTreeSingleIndexDynamicAdaptor_<
2159  Distance, DatasetAdaptor, DIM, AccessorType>;
2160  std::vector<my_kd_tree_t> index_(
2161  treeCount,
2162  my_kd_tree_t(dim /*dim*/, dataset, treeIndex, index_params));
2163  index = index_;
2164  }
2165 
2166  public:
2167  Distance distance;
2168 
2185  const int dimensionality, const DatasetAdaptor& inputData,
2186  const KDTreeSingleIndexAdaptorParams& params =
2188  const size_t maximumPointCount = 1000000000U)
2189  : dataset(inputData), index_params(params), distance(inputData)
2190  {
2191  treeCount = static_cast<size_t>(std::log2(maximumPointCount)) + 1;
2192  pointCount = 0U;
2193  dim = dimensionality;
2194  treeIndex.clear();
2195  if (DIM > 0) dim = DIM;
2196  m_leaf_max_size = params.leaf_max_size;
2197  init();
2198  const size_t num_initial_points = dataset.kdtree_get_point_count();
2199  if (num_initial_points > 0) { addPoints(0, num_initial_points - 1); }
2200  }
2201 
2205  Distance, DatasetAdaptor, DIM, AccessorType>&) = delete;
2206 
2208  void addPoints(AccessorType start, AccessorType end)
2209  {
2210  const Size count = end - start + 1;
2211  int maxIndex = 0;
2212  treeIndex.resize(treeIndex.size() + count);
2213  for (AccessorType idx = start; idx <= end; idx++)
2214  {
2215  const int pos = First0Bit(pointCount);
2216  maxIndex = std::max(pos, maxIndex);
2217  treeIndex[pointCount] = pos;
2218 
2219  const auto it = removedPoints.find(idx);
2220  if (it != removedPoints.end())
2221  {
2222  removedPoints.erase(it);
2223  treeIndex[idx] = pos;
2224  }
2225 
2226  for (int i = 0; i < pos; i++)
2227  {
2228  for (int j = 0; j < static_cast<int>(index[i].vAcc.size()); j++)
2229  {
2230  index[pos].vAcc.push_back(index[i].vAcc[j]);
2231  if (treeIndex[index[i].vAcc[j]] != -1)
2232  treeIndex[index[i].vAcc[j]] = pos;
2233  }
2234  index[i].vAcc.clear();
2235  }
2236  index[pos].vAcc.push_back(idx);
2237  pointCount++;
2238  }
2239 
2240  for (int i = 0; i <= maxIndex; ++i)
2241  {
2242  index[i].freeIndex(index[i]);
2243  if (!index[i].vAcc.empty()) index[i].buildIndex();
2244  }
2245  }
2246 
2248  void removePoint(size_t idx)
2249  {
2250  if (idx >= pointCount) return;
2251  removedPoints.insert(idx);
2252  treeIndex[idx] = -1;
2253  }
2254 
2268  template <typename RESULTSET>
2270  RESULTSET& result, const ElementType* vec,
2271  const SearchParams& searchParams) const
2272  {
2273  for (size_t i = 0; i < treeCount; i++)
2274  {
2275  index[i].findNeighbors(result, &vec[0], searchParams);
2276  }
2277  return result.full();
2278  }
2279 };
2280 
2305 template <
2306  class MatrixType, int32_t DIM = -1, class Distance = nanoflann::metric_L2,
2307  bool row_major = true>
2309 {
2310  using self_t =
2312  using num_t = typename MatrixType::Scalar;
2313  using IndexType = typename MatrixType::Index;
2314  using metric_t = typename Distance::template traits<
2315  num_t, self_t, IndexType>::distance_t;
2316 
2318  metric_t, self_t,
2319  row_major ? MatrixType::ColsAtCompileTime
2320  : MatrixType::RowsAtCompileTime,
2321  IndexType>;
2322 
2323  index_t* index;
2325 
2326  using Offset = typename index_t::Offset;
2327  using Size = typename index_t::Size;
2328  using Dimension = typename index_t::Dimension;
2329 
2332  const Dimension dimensionality,
2333  const std::reference_wrapper<const MatrixType>& mat,
2334  const int leaf_max_size = 10)
2335  : m_data_matrix(mat)
2336  {
2337  const auto dims = row_major ? mat.get().cols() : mat.get().rows();
2338  if (static_cast<Dimension>(dims) != dimensionality)
2339  throw std::runtime_error(
2340  "Error: 'dimensionality' must match column count in data "
2341  "matrix");
2342  if (DIM > 0 && static_cast<int32_t>(dims) != DIM)
2343  throw std::runtime_error(
2344  "Data set dimensionality does not match the 'DIM' template "
2345  "argument");
2346  index = new index_t(
2347  dims, *this /* adaptor */,
2349  }
2350 
2351  public:
2354 
2355  ~KDTreeEigenMatrixAdaptor() { delete index; }
2356 
2357  const std::reference_wrapper<const MatrixType> m_data_matrix;
2358 
2365  inline void query(
2366  const num_t* query_point, const Size num_closest,
2367  IndexType* out_indices, num_t* out_distances_sq,
2368  const int /* nChecks_IGNORED */ = 10) const
2369  {
2370  nanoflann::KNNResultSet<num_t, IndexType> resultSet(num_closest);
2371  resultSet.init(out_indices, out_distances_sq);
2372  index->findNeighbors(resultSet, query_point, nanoflann::SearchParams());
2373  }
2374 
2378  const self_t& derived() const { return *this; }
2379  self_t& derived() { return *this; }
2380 
2381  // Must return the number of data points
2382  inline Size kdtree_get_point_count() const
2383  {
2384  if (row_major)
2385  return m_data_matrix.get().rows();
2386  else
2387  return m_data_matrix.get().cols();
2388  }
2389 
2390  // Returns the dim'th component of the idx'th point in the class:
2391  inline num_t kdtree_get_pt(const IndexType idx, size_t dim) const
2392  {
2393  if (row_major)
2394  return m_data_matrix.get().coeff(idx, IndexType(dim));
2395  else
2396  return m_data_matrix.get().coeff(IndexType(dim), idx);
2397  }
2398 
2399  // Optional bounding-box computation: return false to default to a standard
2400  // bbox computation loop.
2401  // Return true if the BBOX was already computed by the class and returned
2402  // in "bb" so it can be avoided to redo it again. Look at bb.size() to
2403  // find out the expected dimensionality (e.g. 2 or 3 for point clouds)
2404  template <class BBOX>
2405  bool kdtree_get_bbox(BBOX& /*bb*/) const
2406  {
2407  return false;
2408  }
2409 
2412 }; // end of KDTreeEigenMatrixAdaptor // end of grouping
2416 } // namespace nanoflann
Definition: nanoflann.hpp:892
Size veclen(const Derived &obj)
Definition: nanoflann.hpp:981
NodePtr divideTree(Derived &obj, const Offset left, const Offset right, BoundingBox &bbox)
Definition: nanoflann.hpp:1022
PooledAllocator pool
Definition: nanoflann.hpp:975
Dimension dim
Dimensionality of each data point.
Definition: nanoflann.hpp:952
Size m_size_at_index_build
Definition: nanoflann.hpp:950
ElementType dataset_get(const Derived &obj, AccessorType element, Dimension component) const
Helper accessor to the dataset points:
Definition: nanoflann.hpp:984
Size size(const Derived &obj) const
Definition: nanoflann.hpp:978
BoundingBox root_bbox
Definition: nanoflann.hpp:966
Size usedMemory(Derived &obj)
Definition: nanoflann.hpp:994
void saveIndex_(Derived &obj, std::ostream &stream)
Definition: nanoflann.hpp:1229
std::vector< AccessorType > vAcc
Definition: nanoflann.hpp:909
void planeSplit(Derived &obj, Offset ind, const Size count, Dimension cutfeat, DistanceType &cutval, Offset &lim1, Offset &lim2)
Definition: nanoflann.hpp:1141
void freeIndex(Derived &obj)
Definition: nanoflann.hpp:896
void loadIndex_(Derived &obj, std::istream &stream)
Definition: nanoflann.hpp:1244
typename array_or_vector_selector< DIM, Interval >::container_t BoundingBox
Definition: nanoflann.hpp:957
Size m_size
Number of current points in the dataset.
Definition: nanoflann.hpp:949
typename array_or_vector_selector< DIM, DistanceType >::container_t distance_vector_t
Definition: nanoflann.hpp:962
Definition: nanoflann.hpp:1303
KDTreeSingleIndexAdaptor(const Dimension dimensionality, const DatasetAdaptor &inputData, const KDTreeSingleIndexAdaptorParams &params, Args &&... args)
Definition: nanoflann.hpp:1365
void saveIndex(std::ostream &stream)
Definition: nanoflann.hpp:1675
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< std::pair< AccessorType, DistanceType >> &IndicesDists, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1500
void buildIndex()
Definition: nanoflann.hpp:1406
bool searchLevel(RESULTSET &result_set, const ElementType *vec, const NodePtr node, DistanceType mindistsq, distance_vector_t &dists, const float epsError) const
Definition: nanoflann.hpp:1588
KDTreeSingleIndexAdaptor(const KDTreeSingleIndexAdaptor< Distance, DatasetAdaptor, DIM, AccessorType > &)=delete
Size radiusSearchCustomCallback(const ElementType *query_point, SEARCH_CALLBACK &resultSet, const SearchParams &searchParams=SearchParams()) const
Definition: nanoflann.hpp:1521
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1437
typename BaseClassRef::distance_vector_t distance_vector_t
Definition: nanoflann.hpp:1342
Size knnSearch(const ElementType *query_point, const Size num_closest, AccessorType *out_indices, DistanceType *out_distances_sq, const int=10) const
Definition: nanoflann.hpp:1472
void init_vind()
Definition: nanoflann.hpp:1534
void loadIndex(std::istream &stream)
Definition: nanoflann.hpp:1682
const DatasetAdaptor & dataset
The source of our data.
Definition: nanoflann.hpp:1313
typename BaseClassRef::BoundingBox BoundingBox
Definition: nanoflann.hpp:1338
Definition: nanoflann.hpp:1730
const DatasetAdaptor & dataset
The source of our data.
Definition: nanoflann.hpp:1735
void searchLevel(RESULTSET &result_set, const ElementType *vec, const NodePtr node, DistanceType mindistsq, distance_vector_t &dists, const float epsError) const
Definition: nanoflann.hpp:1991
Size radiusSearchCustomCallback(const ElementType *query_point, SEARCH_CALLBACK &resultSet, const SearchParams &searchParams=SearchParams()) const
Definition: nanoflann.hpp:1937
typename BaseClassRef::distance_vector_t distance_vector_t
Definition: nanoflann.hpp:1765
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< std::pair< AccessorType, DistanceType >> &IndicesDists, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1916
void buildIndex()
Definition: nanoflann.hpp:1827
void loadIndex(std::istream &stream)
Definition: nanoflann.hpp:2077
KDTreeSingleIndexDynamicAdaptor_(const KDTreeSingleIndexDynamicAdaptor_ &rhs)=default
typename BaseClassRef::BoundingBox BoundingBox
Definition: nanoflann.hpp:1761
KDTreeSingleIndexDynamicAdaptor_ operator=(const KDTreeSingleIndexDynamicAdaptor_ &rhs)
Definition: nanoflann.hpp:1805
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1856
KDTreeSingleIndexDynamicAdaptor_(const Dimension dimensionality, const DatasetAdaptor &inputData, std::vector< int > &treeIndex_, const KDTreeSingleIndexAdaptorParams &params=KDTreeSingleIndexAdaptorParams())
Definition: nanoflann.hpp:1782
Size knnSearch(const ElementType *query_point, const Size num_closest, AccessorType *out_indices, DistanceType *out_distances_sq, const int=10) const
Definition: nanoflann.hpp:1888
void saveIndex(std::ostream &stream)
Definition: nanoflann.hpp:2070
Definition: nanoflann.hpp:2098
void removePoint(size_t idx)
Definition: nanoflann.hpp:2248
const DatasetAdaptor & dataset
The source of our data.
Definition: nanoflann.hpp:2118
std::vector< int > treeIndex
Definition: nanoflann.hpp:2121
void addPoints(AccessorType start, AccessorType end)
Definition: nanoflann.hpp:2208
Dimension dim
Dimensionality of each data point.
Definition: nanoflann.hpp:2128
KDTreeSingleIndexDynamicAdaptor(const KDTreeSingleIndexDynamicAdaptor< Distance, DatasetAdaptor, DIM, AccessorType > &)=delete
const std::vector< index_container_t > & getAllIndices() const
Definition: nanoflann.hpp:2137
KDTreeSingleIndexDynamicAdaptor(const int dimensionality, const DatasetAdaptor &inputData, const KDTreeSingleIndexAdaptorParams &params=KDTreeSingleIndexAdaptorParams(), const size_t maximumPointCount=1000000000U)
Definition: nanoflann.hpp:2184
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParams &searchParams) const
Definition: nanoflann.hpp:2269
Definition: nanoflann.hpp:159
bool addPoint(DistanceType dist, IndexType index)
Definition: nanoflann.hpp:195
Definition: nanoflann.hpp:736
T * allocate(const size_t count=1)
Definition: nanoflann.hpp:847
~PooledAllocator()
Definition: nanoflann.hpp:771
void free_all()
Definition: nanoflann.hpp:774
PooledAllocator()
Definition: nanoflann.hpp:766
void * malloc(const size_t req_size)
Definition: nanoflann.hpp:790
Definition: nanoflann.hpp:249
std::pair< IndexType, DistanceType > worst_item() const
Definition: nanoflann.hpp:292
bool addPoint(DistanceType dist, IndexType index)
Definition: nanoflann.hpp:279
const size_t WORDSIZE
Definition: nanoflann.hpp:732
T * allocate(size_t count=1)
Definition: nanoflann.hpp:711
T pi_const()
Definition: nanoflann.hpp:80
std::enable_if< has_assign< Container >::value, void >::type assign(Container &c, const size_t nElements, const T &value)
Definition: nanoflann.hpp:137
std::enable_if< has_resize< Container >::value, void >::type resize(Container &c, const size_t nElements)
Definition: nanoflann.hpp:115
Definition: nanoflann.hpp:235
bool operator()(const PairType &p1, const PairType &p2) const
Definition: nanoflann.hpp:238
Definition: nanoflann.hpp:941
Definition: nanoflann.hpp:918
DistanceType divhigh
The values used for subdivision.
Definition: nanoflann.hpp:931
Dimension divfeat
Dimension used for subdivision.
Definition: nanoflann.hpp:929
Offset right
Indices of points in leaf node.
Definition: nanoflann.hpp:925
union nanoflann::KDTreeBaseClass::Node::@0 node_type
Node * child1
Definition: nanoflann.hpp:935
Definition: nanoflann.hpp:2309
void query(const num_t *query_point, const Size num_closest, IndexType *out_indices, num_t *out_distances_sq, const int=10) const
Definition: nanoflann.hpp:2365
KDTreeEigenMatrixAdaptor(const Dimension dimensionality, const std::reference_wrapper< const MatrixType > &mat, const int leaf_max_size=10)
Constructor: takes a const ref to the matrix object with the data points.
Definition: nanoflann.hpp:2331
KDTreeEigenMatrixAdaptor(const self_t &)=delete
typename index_t::Offset Offset
Definition: nanoflann.hpp:2326
Definition: nanoflann.hpp:670
Definition: nanoflann.hpp:361
Definition: nanoflann.hpp:423
Definition: nanoflann.hpp:488
Definition: nanoflann.hpp:344
Definition: nanoflann.hpp:533
DistanceType accum_dist(const U a, const V b, const size_t) const
Definition: nanoflann.hpp:551
Definition: nanoflann.hpp:578
Definition: nanoflann.hpp:684
bool sorted
distance (default: true)
Definition: nanoflann.hpp:695
SearchParams(int checks_IGNORED_=32, float eps_=0, bool sorted_=true)
Definition: nanoflann.hpp:687
float eps
search for eps-approximate neighbours (default: 0)
Definition: nanoflann.hpp:694
int checks
Definition: nanoflann.hpp:692
Definition: nanoflann.hpp:863
Definition: nanoflann.hpp:102
Definition: nanoflann.hpp:91
Definition: nanoflann.hpp:608
Definition: nanoflann.hpp:605
Definition: nanoflann.hpp:617
Definition: nanoflann.hpp:626
Definition: nanoflann.hpp:623
Definition: nanoflann.hpp:614
Definition: nanoflann.hpp:635
Definition: nanoflann.hpp:632
Definition: nanoflann.hpp:644
Definition: nanoflann.hpp:641