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 // Avoid conflicts with X11 headers
73 #ifdef None
74 #undef None
75 #endif
76 
77 namespace nanoflann
78 {
83 template <typename T>
85 {
86  return static_cast<T>(3.14159265358979323846);
87 }
88 
93 template <typename T, typename = int>
94 struct has_resize : std::false_type
95 {
96 };
97 
98 template <typename T>
99 struct has_resize<T, decltype((void)std::declval<T>().resize(1), 0)>
100  : std::true_type
101 {
102 };
103 
104 template <typename T, typename = int>
105 struct has_assign : std::false_type
106 {
107 };
108 
109 template <typename T>
110 struct has_assign<T, decltype((void)std::declval<T>().assign(1, 0), 0)>
111  : std::true_type
112 {
113 };
114 
118 template <typename Container>
119 inline typename std::enable_if<has_resize<Container>::value, void>::type resize(
120  Container& c, const size_t nElements)
121 {
122  c.resize(nElements);
123 }
124 
129 template <typename Container>
130 inline typename std::enable_if<!has_resize<Container>::value, void>::type
131  resize(Container& c, const size_t nElements)
132 {
133  if (nElements != c.size())
134  throw std::logic_error("Try to change the size of a std::array.");
135 }
136 
140 template <typename Container, typename T>
141 inline typename std::enable_if<has_assign<Container>::value, void>::type assign(
142  Container& c, const size_t nElements, const T& value)
143 {
144  c.assign(nElements, value);
145 }
146 
150 template <typename Container, typename T>
151 inline typename std::enable_if<!has_assign<Container>::value, void>::type
152  assign(Container& c, const size_t nElements, const T& value)
153 {
154  for (size_t i = 0; i < nElements; i++) c[i] = value;
155 }
156 
159 template <
160  typename _DistanceType, typename _IndexType = size_t,
161  typename _CountType = size_t>
163 {
164  public:
165  using DistanceType = _DistanceType;
166  using IndexType = _IndexType;
167  using CountType = _CountType;
168 
169  private:
170  IndexType* indices;
171  DistanceType* dists;
172  CountType capacity;
173  CountType count;
174 
175  public:
176  explicit inline KNNResultSet(CountType capacity_)
177  : indices(0), dists(0), capacity(capacity_), count(0)
178  {
179  }
180 
181  inline void init(IndexType* indices_, DistanceType* dists_)
182  {
183  indices = indices_;
184  dists = dists_;
185  count = 0;
186  if (capacity)
187  dists[capacity - 1] = (std::numeric_limits<DistanceType>::max)();
188  }
189 
190  inline CountType size() const { return count; }
191 
192  inline bool full() const { return count == capacity; }
193 
199  inline bool addPoint(DistanceType dist, IndexType index)
200  {
201  CountType i;
202  for (i = count; i > 0; --i)
203  {
204 #ifdef NANOFLANN_FIRST_MATCH // If defined and two points have the same
205  // distance, the one with the lowest-index will be
206  // returned first.
207  if ((dists[i - 1] > dist) ||
208  ((dist == dists[i - 1]) && (indices[i - 1] > index)))
209  {
210 #else
211  if (dists[i - 1] > dist)
212  {
213 #endif
214  if (i < capacity)
215  {
216  dists[i] = dists[i - 1];
217  indices[i] = indices[i - 1];
218  }
219  }
220  else
221  break;
222  }
223  if (i < capacity)
224  {
225  dists[i] = dist;
226  indices[i] = index;
227  }
228  if (count < capacity) count++;
229 
230  // tell caller that the search shall continue
231  return true;
232  }
233 
234  inline DistanceType worstDist() const { return dists[capacity - 1]; }
235 };
236 
239 {
241  template <typename PairType>
242  inline bool operator()(const PairType& p1, const PairType& p2) const
243  {
244  return p1.second < p2.second;
245  }
246 };
247 
251 template <typename _DistanceType, typename _IndexType = size_t>
253 {
254  public:
255  using DistanceType = _DistanceType;
256  using IndexType = _IndexType;
257 
258  public:
259  const DistanceType radius;
260 
261  std::vector<std::pair<IndexType, DistanceType>>& m_indices_dists;
262 
263  explicit inline RadiusResultSet(
264  DistanceType radius_,
265  std::vector<std::pair<IndexType, DistanceType>>& indices_dists)
266  : radius(radius_), m_indices_dists(indices_dists)
267  {
268  init();
269  }
270 
271  inline void init() { clear(); }
272  inline void clear() { m_indices_dists.clear(); }
273 
274  inline size_t size() const { return m_indices_dists.size(); }
275 
276  inline bool full() const { return true; }
277 
283  inline bool addPoint(DistanceType dist, IndexType index)
284  {
285  if (dist < radius)
286  m_indices_dists.push_back(std::make_pair(index, dist));
287  return true;
288  }
289 
290  inline DistanceType worstDist() const { return radius; }
291 
296  std::pair<IndexType, DistanceType> worst_item() const
297  {
298  if (m_indices_dists.empty())
299  throw std::runtime_error(
300  "Cannot invoke RadiusResultSet::worst_item() on "
301  "an empty list of results.");
302  using DistIt = typename std::vector<
303  std::pair<IndexType, DistanceType>>::const_iterator;
304  DistIt it = std::max_element(
305  m_indices_dists.begin(), m_indices_dists.end(), IndexDist_Sorter());
306  return *it;
307  }
308 };
309 
314 template <typename T>
315 void save_value(std::ostream& stream, const T& value)
316 {
317  stream.write(reinterpret_cast<const char*>(&value), sizeof(T));
318 }
319 
320 template <typename T>
321 void save_value(std::ostream& stream, const std::vector<T>& value)
322 {
323  size_t size = value.size();
324  stream.write(reinterpret_cast<const char*>(&size), sizeof(size_t));
325  stream.write(reinterpret_cast<const char*>(value.data()), sizeof(T) * size);
326 }
327 
328 template <typename T>
329 void load_value(std::istream& stream, T& value)
330 {
331  stream.read(reinterpret_cast<char*>(&value), sizeof(T));
332 }
333 
334 template <typename T>
335 void load_value(std::istream& stream, std::vector<T>& value)
336 {
337  size_t size;
338  stream.read(reinterpret_cast<char*>(&size), sizeof(size_t));
339  value.resize(size);
340  stream.read(reinterpret_cast<char*>(value.data()), sizeof(T) * size);
341 }
347 struct Metric
348 {
349 };
350 
361 template <
362  class T, class DataSource, typename _DistanceType = T,
363  typename AccessorType = uint32_t>
365 {
366  using ElementType = T;
367  using DistanceType = _DistanceType;
368 
369  const DataSource& data_source;
370 
371  L1_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
372 
373  inline DistanceType evalMetric(
374  const T* a, const AccessorType b_idx, size_t size,
375  DistanceType worst_dist = -1) const
376  {
377  DistanceType result = DistanceType();
378  const T* last = a + size;
379  const T* lastgroup = last - 3;
380  size_t d = 0;
381 
382  /* Process 4 items with each loop for efficiency. */
383  while (a < lastgroup)
384  {
385  const DistanceType diff0 =
386  std::abs(a[0] - data_source.kdtree_get_pt(b_idx, d++));
387  const DistanceType diff1 =
388  std::abs(a[1] - data_source.kdtree_get_pt(b_idx, d++));
389  const DistanceType diff2 =
390  std::abs(a[2] - data_source.kdtree_get_pt(b_idx, d++));
391  const DistanceType diff3 =
392  std::abs(a[3] - data_source.kdtree_get_pt(b_idx, d++));
393  result += diff0 + diff1 + diff2 + diff3;
394  a += 4;
395  if ((worst_dist > 0) && (result > worst_dist)) { return result; }
396  }
397  /* Process last 0-3 components. Not needed for standard vector lengths.
398  */
399  while (a < last)
400  {
401  result += std::abs(*a++ - data_source.kdtree_get_pt(b_idx, d++));
402  }
403  return result;
404  }
405 
406  template <typename U, typename V>
407  inline DistanceType accum_dist(const U a, const V b, const size_t) const
408  {
409  return std::abs(a - b);
410  }
411 };
412 
423 template <
424  class T, class DataSource, typename _DistanceType = T,
425  typename AccessorType = uint32_t>
427 {
428  using ElementType = T;
429  using DistanceType = _DistanceType;
430 
431  const DataSource& data_source;
432 
433  L2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
434 
435  inline DistanceType evalMetric(
436  const T* a, const AccessorType b_idx, size_t size,
437  DistanceType worst_dist = -1) const
438  {
439  DistanceType result = DistanceType();
440  const T* last = a + size;
441  const T* lastgroup = last - 3;
442  size_t d = 0;
443 
444  /* Process 4 items with each loop for efficiency. */
445  while (a < lastgroup)
446  {
447  const DistanceType diff0 =
448  a[0] - data_source.kdtree_get_pt(b_idx, d++);
449  const DistanceType diff1 =
450  a[1] - data_source.kdtree_get_pt(b_idx, d++);
451  const DistanceType diff2 =
452  a[2] - data_source.kdtree_get_pt(b_idx, d++);
453  const DistanceType diff3 =
454  a[3] - data_source.kdtree_get_pt(b_idx, d++);
455  result +=
456  diff0 * diff0 + diff1 * diff1 + diff2 * diff2 + diff3 * diff3;
457  a += 4;
458  if ((worst_dist > 0) && (result > worst_dist)) { return result; }
459  }
460  /* Process last 0-3 components. Not needed for standard vector lengths.
461  */
462  while (a < last)
463  {
464  const DistanceType diff0 =
465  *a++ - data_source.kdtree_get_pt(b_idx, d++);
466  result += diff0 * diff0;
467  }
468  return result;
469  }
470 
471  template <typename U, typename V>
472  inline DistanceType accum_dist(const U a, const V b, const size_t) const
473  {
474  return (a - b) * (a - b);
475  }
476 };
477 
488 template <
489  class T, class DataSource, typename _DistanceType = T,
490  typename AccessorType = uint32_t>
492 {
493  using ElementType = T;
494  using DistanceType = _DistanceType;
495 
496  const DataSource& data_source;
497 
498  L2_Simple_Adaptor(const DataSource& _data_source)
499  : data_source(_data_source)
500  {
501  }
502 
503  inline DistanceType evalMetric(
504  const T* a, const AccessorType b_idx, size_t size) const
505  {
506  DistanceType result = DistanceType();
507  for (size_t i = 0; i < size; ++i)
508  {
509  const DistanceType diff =
510  a[i] - data_source.kdtree_get_pt(b_idx, i);
511  result += diff * diff;
512  }
513  return result;
514  }
515 
516  template <typename U, typename V>
517  inline DistanceType accum_dist(const U a, const V b, const size_t) const
518  {
519  return (a - b) * (a - b);
520  }
521 };
522 
533 template <
534  class T, class DataSource, typename _DistanceType = T,
535  typename AccessorType = uint32_t>
537 {
538  using ElementType = T;
539  using DistanceType = _DistanceType;
540 
541  const DataSource& data_source;
542 
543  SO2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
544 
545  inline DistanceType evalMetric(
546  const T* a, const AccessorType b_idx, size_t size) const
547  {
548  return accum_dist(
549  a[size - 1], data_source.kdtree_get_pt(b_idx, size - 1), size - 1);
550  }
551 
554  template <typename U, typename V>
555  inline DistanceType accum_dist(const U a, const V b, const size_t) const
556  {
557  DistanceType result = DistanceType();
558  DistanceType PI = pi_const<DistanceType>();
559  result = b - a;
560  if (result > PI)
561  result -= 2 * PI;
562  else if (result < -PI)
563  result += 2 * PI;
564  return result;
565  }
566 };
567 
578 template <
579  class T, class DataSource, typename _DistanceType = T,
580  typename AccessorType = uint32_t>
582 {
583  using ElementType = T;
584  using DistanceType = _DistanceType;
585 
587  distance_L2_Simple;
588 
589  SO3_Adaptor(const DataSource& _data_source)
590  : distance_L2_Simple(_data_source)
591  {
592  }
593 
594  inline DistanceType evalMetric(
595  const T* a, const AccessorType b_idx, size_t size) const
596  {
597  return distance_L2_Simple.evalMetric(a, b_idx, size);
598  }
599 
600  template <typename U, typename V>
601  inline DistanceType accum_dist(const U a, const V b, const size_t idx) const
602  {
603  return distance_L2_Simple.accum_dist(a, b, idx);
604  }
605 };
606 
608 struct metric_L1 : public Metric
609 {
610  template <class T, class DataSource, typename AccessorType = uint32_t>
611  struct traits
612  {
614  };
615 };
617 struct metric_L2 : public Metric
618 {
619  template <class T, class DataSource, typename AccessorType = uint32_t>
620  struct traits
621  {
623  };
624 };
626 struct metric_L2_Simple : public Metric
627 {
628  template <class T, class DataSource, typename AccessorType = uint32_t>
629  struct traits
630  {
632  };
633 };
635 struct metric_SO2 : public Metric
636 {
637  template <class T, class DataSource, typename AccessorType = uint32_t>
638  struct traits
639  {
641  };
642 };
644 struct metric_SO3 : public Metric
645 {
646  template <class T, class DataSource, typename AccessorType = uint32_t>
647  struct traits
648  {
650  };
651 };
652 
658 enum class KDTreeSingleIndexAdaptorFlags
659 {
660  None = 0,
661  SkipInitialBuildIndex = 1
662 };
663 
664 inline std::underlying_type<KDTreeSingleIndexAdaptorFlags>::type operator&(
665  KDTreeSingleIndexAdaptorFlags lhs, KDTreeSingleIndexAdaptorFlags rhs)
666 {
667  using underlying =
668  typename std::underlying_type<KDTreeSingleIndexAdaptorFlags>::type;
669  return static_cast<underlying>(lhs) & static_cast<underlying>(rhs);
670 }
671 
674 {
676  size_t _leaf_max_size = 10, KDTreeSingleIndexAdaptorFlags _flags =
677  KDTreeSingleIndexAdaptorFlags::None)
678  : leaf_max_size(_leaf_max_size), flags(_flags)
679  {
680  }
681 
682  size_t leaf_max_size;
683  KDTreeSingleIndexAdaptorFlags flags;
684 };
685 
688 {
691  SearchParams(int checks_IGNORED_ = 32, float eps_ = 0, bool sorted_ = true)
692  : checks(checks_IGNORED_), eps(eps_), sorted(sorted_)
693  {
694  }
695 
696  int checks;
698  float eps;
699  bool sorted;
701 };
714 template <typename T>
715 inline T* allocate(size_t count = 1)
716 {
717  T* mem = static_cast<T*>(::malloc(sizeof(T) * count));
718  return mem;
719 }
720 
736 const size_t WORDSIZE = 16;
737 const size_t BLOCKSIZE = 8192;
738 
740 {
741  /* We maintain memory alignment to word boundaries by requiring that all
742  allocations be in multiples of the machine wordsize. */
743  /* Size of machine word in bytes. Must be power of 2. */
744  /* Minimum number of bytes requested at a time from the system. Must be
745  * multiple of WORDSIZE. */
746 
747  using Offset = uint32_t;
748  using Size = uint32_t;
749  using Dimension = int32_t;
750 
751  Size remaining; /* Number of bytes left in current block of storage. */
752  void* base; /* Pointer to base of current block of storage. */
753  void* loc; /* Current location in block to next allocate memory. */
754 
755  void internal_init()
756  {
757  remaining = 0;
758  base = nullptr;
759  usedMemory = 0;
760  wastedMemory = 0;
761  }
762 
763  public:
764  Size usedMemory;
765  Size wastedMemory;
766 
770  PooledAllocator() { internal_init(); }
771 
775  ~PooledAllocator() { free_all(); }
776 
778  void free_all()
779  {
780  while (base != nullptr)
781  {
782  void* prev =
783  *(static_cast<void**>(base)); /* Get pointer to prev block. */
784  ::free(base);
785  base = prev;
786  }
787  internal_init();
788  }
789 
794  void* malloc(const size_t req_size)
795  {
796  /* Round size up to a multiple of wordsize. The following expression
797  only works for WORDSIZE that is a power of 2, by masking last bits
798  of incremented size to zero.
799  */
800  const Size size = (req_size + (WORDSIZE - 1)) & ~(WORDSIZE - 1);
801 
802  /* Check whether a new block must be allocated. Note that the first
803  word of a block is reserved for a pointer to the previous block.
804  */
805  if (size > remaining)
806  {
807  wastedMemory += remaining;
808 
809  /* Allocate new storage. */
810  const Size blocksize =
811  (size + sizeof(void*) + (WORDSIZE - 1) > BLOCKSIZE)
812  ? size + sizeof(void*) + (WORDSIZE - 1)
813  : BLOCKSIZE;
814 
815  // use the standard C malloc to allocate memory
816  void* m = ::malloc(blocksize);
817  if (!m)
818  {
819  fprintf(stderr, "Failed to allocate memory.\n");
820  throw std::bad_alloc();
821  }
822 
823  /* Fill first word of new block with pointer to previous block. */
824  static_cast<void**>(m)[0] = base;
825  base = m;
826 
827  Size shift = 0;
828  // int size_t = (WORDSIZE - ( (((size_t)m) + sizeof(void*)) &
829  // (WORDSIZE-1))) & (WORDSIZE-1);
830 
831  remaining = blocksize - sizeof(void*) - shift;
832  loc = (static_cast<char*>(m) + sizeof(void*) + shift);
833  }
834  void* rloc = loc;
835  loc = static_cast<char*>(loc) + size;
836  remaining -= size;
837 
838  usedMemory += size;
839 
840  return rloc;
841  }
842 
850  template <typename T>
851  T* allocate(const size_t count = 1)
852  {
853  T* mem = static_cast<T*>(this->malloc(sizeof(T) * count));
854  return mem;
855  }
856 };
865 template <int32_t DIM, typename T>
867 {
868  using container_t = std::array<T, DIM>;
869 };
871 template <typename T>
873 {
874  using container_t = std::vector<T>;
875 };
876 
892 template <
893  class Derived, typename Distance, class DatasetAdaptor, int32_t DIM = -1,
894  typename AccessorType = uint32_t>
896 {
897  public:
900  void freeIndex(Derived& obj)
901  {
902  obj.pool.free_all();
903  obj.root_node = nullptr;
904  obj.m_size_at_index_build = 0;
905  }
906 
907  using ElementType = typename Distance::ElementType;
908  using DistanceType = typename Distance::DistanceType;
909 
913  std::vector<AccessorType> vAcc;
914 
915  using Offset = typename decltype(vAcc)::size_type;
916  using Size = typename decltype(vAcc)::size_type;
917  using Dimension = int32_t;
918 
919  /*--------------------- Internal Data Structures
920  * --------------------------*/
921  struct Node
922  {
925  union
926  {
927  struct leaf
928  {
929  Offset left, right;
930  } lr;
931  struct nonleaf
932  {
933  Dimension divfeat;
934  DistanceType divlow,
936  } sub;
937  } node_type;
939  Node *child1, *child2;
940  };
941 
942  using NodePtr = Node*;
943 
944  struct Interval
945  {
946  ElementType low, high;
947  };
948 
949  NodePtr root_node;
950 
951  Size m_leaf_max_size;
952 
953  Size m_size;
956  Dimension dim;
957 
960  using BoundingBox =
962 
966  typename array_or_vector_selector<DIM, DistanceType>::container_t;
967 
971 
980 
982  Size size(const Derived& obj) const { return obj.m_size; }
983 
985  Size veclen(const Derived& obj) { return DIM > 0 ? DIM : obj.dim; }
986 
988  inline ElementType dataset_get(
989  const Derived& obj, AccessorType element, Dimension component) const
990  {
991  return obj.dataset.kdtree_get_pt(element, component);
992  }
993 
998  Size usedMemory(Derived& obj)
999  {
1000  return obj.pool.usedMemory + obj.pool.wastedMemory +
1001  obj.dataset.kdtree_get_point_count() *
1002  sizeof(AccessorType); // pool memory and vind array memory
1003  }
1004 
1005  void computeMinMax(
1006  const Derived& obj, Offset ind, Size count, Dimension element,
1007  ElementType& min_elem, ElementType& max_elem)
1008  {
1009  min_elem = dataset_get(obj, vAcc[ind], element);
1010  max_elem = min_elem;
1011  for (Offset i = 1; i < count; ++i)
1012  {
1013  ElementType val = dataset_get(obj, vAcc[ind + i], element);
1014  if (val < min_elem) min_elem = val;
1015  if (val > max_elem) max_elem = val;
1016  }
1017  }
1018 
1027  Derived& obj, const Offset left, const Offset right, BoundingBox& bbox)
1028  {
1029  NodePtr node = obj.pool.template allocate<Node>(); // allocate memory
1030 
1031  /* If too few exemplars remain, then make this a leaf node. */
1032  if ((right - left) <= static_cast<Offset>(obj.m_leaf_max_size))
1033  {
1034  node->child1 = node->child2 = nullptr; /* Mark as leaf node. */
1035  node->node_type.lr.left = left;
1036  node->node_type.lr.right = right;
1037 
1038  // compute bounding-box of leaf points
1039  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1040  {
1041  bbox[i].low = dataset_get(obj, obj.vAcc[left], i);
1042  bbox[i].high = dataset_get(obj, obj.vAcc[left], i);
1043  }
1044  for (Offset k = left + 1; k < right; ++k)
1045  {
1046  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1047  {
1048  if (bbox[i].low > dataset_get(obj, obj.vAcc[k], i))
1049  bbox[i].low = dataset_get(obj, obj.vAcc[k], i);
1050  if (bbox[i].high < dataset_get(obj, obj.vAcc[k], i))
1051  bbox[i].high = dataset_get(obj, obj.vAcc[k], i);
1052  }
1053  }
1054  }
1055  else
1056  {
1057  Offset idx;
1058  Dimension cutfeat;
1059  DistanceType cutval;
1060  middleSplit_(obj, left, right - left, idx, cutfeat, cutval, bbox);
1061 
1062  node->node_type.sub.divfeat = cutfeat;
1063 
1064  BoundingBox left_bbox(bbox);
1065  left_bbox[cutfeat].high = cutval;
1066  node->child1 = divideTree(obj, left, left + idx, left_bbox);
1067 
1068  BoundingBox right_bbox(bbox);
1069  right_bbox[cutfeat].low = cutval;
1070  node->child2 = divideTree(obj, left + idx, right, right_bbox);
1071 
1072  node->node_type.sub.divlow = left_bbox[cutfeat].high;
1073  node->node_type.sub.divhigh = right_bbox[cutfeat].low;
1074 
1075  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1076  {
1077  bbox[i].low = std::min(left_bbox[i].low, right_bbox[i].low);
1078  bbox[i].high = std::max(left_bbox[i].high, right_bbox[i].high);
1079  }
1080  }
1081 
1082  return node;
1083  }
1084 
1085  void middleSplit_(
1086  Derived& obj, Offset ind, Size count, Offset& index, Dimension& cutfeat,
1087  DistanceType& cutval, const BoundingBox& bbox)
1088  {
1089  const auto EPS = static_cast<DistanceType>(0.00001);
1090  ElementType max_span = bbox[0].high - bbox[0].low;
1091  for (Dimension i = 1; i < (DIM > 0 ? DIM : obj.dim); ++i)
1092  {
1093  ElementType span = bbox[i].high - bbox[i].low;
1094  if (span > max_span) { max_span = span; }
1095  }
1096  ElementType max_spread = -1;
1097  cutfeat = 0;
1098  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1099  {
1100  ElementType span = bbox[i].high - bbox[i].low;
1101  if (span > (1 - EPS) * max_span)
1102  {
1103  ElementType min_elem, max_elem;
1104  computeMinMax(obj, ind, count, i, min_elem, max_elem);
1105  ElementType spread = max_elem - min_elem;
1106  if (spread > max_spread)
1107  {
1108  cutfeat = i;
1109  max_spread = spread;
1110  }
1111  }
1112  }
1113  // split in the middle
1114  DistanceType split_val = (bbox[cutfeat].low + bbox[cutfeat].high) / 2;
1115  ElementType min_elem, max_elem;
1116  computeMinMax(obj, ind, count, cutfeat, min_elem, max_elem);
1117 
1118  if (split_val < min_elem)
1119  cutval = min_elem;
1120  else if (split_val > max_elem)
1121  cutval = max_elem;
1122  else
1123  cutval = split_val;
1124 
1125  Offset lim1, lim2;
1126  planeSplit(obj, ind, count, cutfeat, cutval, lim1, lim2);
1127 
1128  if (lim1 > count / 2)
1129  index = lim1;
1130  else if (lim2 < count / 2)
1131  index = lim2;
1132  else
1133  index = count / 2;
1134  }
1135 
1146  Derived& obj, Offset ind, const Size count, Dimension cutfeat,
1147  DistanceType& cutval, Offset& lim1, Offset& lim2)
1148  {
1149  /* Move vector indices for left subtree to front of list. */
1150  Offset left = 0;
1151  Offset right = count - 1;
1152  for (;;)
1153  {
1154  while (left <= right &&
1155  dataset_get(obj, vAcc[ind + left], cutfeat) < cutval)
1156  ++left;
1157  while (right && left <= right &&
1158  dataset_get(obj, vAcc[ind + right], cutfeat) >= cutval)
1159  --right;
1160  if (left > right || !right)
1161  break; // "!right" was added to support unsigned Index types
1162  std::swap(vAcc[ind + left], vAcc[ind + right]);
1163  ++left;
1164  --right;
1165  }
1166  /* If either list is empty, it means that all remaining features
1167  * are identical. Split in the middle to maintain a balanced tree.
1168  */
1169  lim1 = left;
1170  right = count - 1;
1171  for (;;)
1172  {
1173  while (left <= right &&
1174  dataset_get(obj, vAcc[ind + left], cutfeat) <= cutval)
1175  ++left;
1176  while (right && left <= right &&
1177  dataset_get(obj, vAcc[ind + right], cutfeat) > cutval)
1178  --right;
1179  if (left > right || !right)
1180  break; // "!right" was added to support unsigned Index types
1181  std::swap(vAcc[ind + left], vAcc[ind + right]);
1182  ++left;
1183  --right;
1184  }
1185  lim2 = left;
1186  }
1187 
1188  DistanceType computeInitialDistances(
1189  const Derived& obj, const ElementType* vec,
1190  distance_vector_t& dists) const
1191  {
1192  assert(vec);
1193  DistanceType distsq = DistanceType();
1194 
1195  for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim); ++i)
1196  {
1197  if (vec[i] < obj.root_bbox[i].low)
1198  {
1199  dists[i] =
1200  obj.distance.accum_dist(vec[i], obj.root_bbox[i].low, i);
1201  distsq += dists[i];
1202  }
1203  if (vec[i] > obj.root_bbox[i].high)
1204  {
1205  dists[i] =
1206  obj.distance.accum_dist(vec[i], obj.root_bbox[i].high, i);
1207  distsq += dists[i];
1208  }
1209  }
1210  return distsq;
1211  }
1212 
1213  void save_tree(Derived& obj, std::ostream& stream, NodePtr tree)
1214  {
1215  save_value(stream, *tree);
1216  if (tree->child1 != nullptr) { save_tree(obj, stream, tree->child1); }
1217  if (tree->child2 != nullptr) { save_tree(obj, stream, tree->child2); }
1218  }
1219 
1220  void load_tree(Derived& obj, std::istream& stream, NodePtr& tree)
1221  {
1222  tree = obj.pool.template allocate<Node>();
1223  load_value(stream, *tree);
1224  if (tree->child1 != nullptr) { load_tree(obj, stream, tree->child1); }
1225  if (tree->child2 != nullptr) { load_tree(obj, stream, tree->child2); }
1226  }
1227 
1233  void saveIndex_(Derived& obj, std::ostream& stream)
1234  {
1235  save_value(stream, obj.m_size);
1236  save_value(stream, obj.dim);
1237  save_value(stream, obj.root_bbox);
1238  save_value(stream, obj.m_leaf_max_size);
1239  save_value(stream, obj.vAcc);
1240  save_tree(obj, stream, obj.root_node);
1241  }
1242 
1248  void loadIndex_(Derived& obj, std::istream& stream)
1249  {
1250  load_value(stream, obj.m_size);
1251  load_value(stream, obj.dim);
1252  load_value(stream, obj.root_bbox);
1253  load_value(stream, obj.m_leaf_max_size);
1254  load_value(stream, obj.vAcc);
1255  load_tree(obj, stream, obj.root_node);
1256  }
1257 };
1258 
1300 template <
1301  typename Distance, class DatasetAdaptor, int32_t DIM = -1,
1302  typename AccessorType = uint32_t>
1304  : public KDTreeBaseClass<
1305  KDTreeSingleIndexAdaptor<Distance, DatasetAdaptor, DIM, AccessorType>,
1306  Distance, DatasetAdaptor, DIM, AccessorType>
1307 {
1308  public:
1312  Distance, DatasetAdaptor, DIM, AccessorType>&) = delete;
1313 
1317  const DatasetAdaptor& dataset;
1318 
1319  const KDTreeSingleIndexAdaptorParams index_params;
1320 
1321  Distance distance;
1322 
1323  using BaseClassRef = typename nanoflann::KDTreeBaseClass<
1325  Distance, DatasetAdaptor, DIM, AccessorType>,
1326  Distance, DatasetAdaptor, DIM, AccessorType>;
1327 
1328  using Offset = typename BaseClassRef::Offset;
1329  using Size = typename BaseClassRef::Size;
1330  using Dimension = typename BaseClassRef::Dimension;
1331 
1332  using ElementType = typename BaseClassRef::ElementType;
1333  using DistanceType = typename BaseClassRef::DistanceType;
1334 
1335  using Node = typename BaseClassRef::Node;
1336  using NodePtr = Node*;
1337 
1338  using Interval = typename BaseClassRef::Interval;
1339 
1342  using BoundingBox = typename BaseClassRef::BoundingBox;
1343 
1346  using distance_vector_t = typename BaseClassRef::distance_vector_t;
1347 
1368  template <class... Args>
1370  const Dimension dimensionality, const DatasetAdaptor& inputData,
1371  const KDTreeSingleIndexAdaptorParams& params, Args&&... args)
1372  : dataset(inputData),
1373  index_params(params),
1374  distance(inputData, std::forward<Args>(args)...)
1375  {
1376  init(dimensionality, params);
1377  }
1378 
1379  explicit KDTreeSingleIndexAdaptor(
1380  const Dimension dimensionality, const DatasetAdaptor& inputData,
1381  const KDTreeSingleIndexAdaptorParams& params = {})
1382  : dataset(inputData), index_params(params), distance(inputData)
1383  {
1384  init(dimensionality, params);
1385  }
1386 
1387  private:
1388  void init(
1389  const Dimension dimensionality,
1390  const KDTreeSingleIndexAdaptorParams& params)
1391  {
1392  BaseClassRef::root_node = nullptr;
1393  BaseClassRef::m_size = dataset.kdtree_get_point_count();
1394  BaseClassRef::m_size_at_index_build = BaseClassRef::m_size;
1395  BaseClassRef::dim = dimensionality;
1396  if (DIM > 0) BaseClassRef::dim = DIM;
1397  BaseClassRef::m_leaf_max_size = params.leaf_max_size;
1398 
1399  if (!(params.flags &
1400  KDTreeSingleIndexAdaptorFlags::SkipInitialBuildIndex))
1401  {
1402  buildIndex();
1403  }
1404  }
1405 
1406  public:
1410  void buildIndex()
1411  {
1412  BaseClassRef::m_size = dataset.kdtree_get_point_count();
1413  BaseClassRef::m_size_at_index_build = BaseClassRef::m_size;
1414  init_vind();
1415  this->freeIndex(*this);
1416  BaseClassRef::m_size_at_index_build = BaseClassRef::m_size;
1417  if (BaseClassRef::m_size == 0) return;
1418  computeBoundingBox(BaseClassRef::root_bbox);
1419  BaseClassRef::root_node = this->divideTree(
1420  *this, 0, BaseClassRef::m_size,
1421  BaseClassRef::root_bbox); // construct the tree
1422  }
1423 
1440  template <typename RESULTSET>
1442  RESULTSET& result, const ElementType* vec,
1443  const SearchParams& searchParams) const
1444  {
1445  assert(vec);
1446  if (this->size(*this) == 0) return false;
1447  if (!BaseClassRef::root_node)
1448  throw std::runtime_error(
1449  "[nanoflann] findNeighbors() called before building the "
1450  "index.");
1451  float epsError = 1 + searchParams.eps;
1452 
1454  dists; // fixed or variable-sized container (depending on DIM)
1455  auto zero = static_cast<decltype(result.worstDist())>(0);
1456  assign(
1457  dists, (DIM > 0 ? DIM : BaseClassRef::dim),
1458  zero); // Fill it with zeros.
1459  DistanceType distsq = this->computeInitialDistances(*this, vec, dists);
1460  searchLevel(
1461  result, vec, BaseClassRef::root_node, distsq, dists,
1462  epsError); // "count_leaf" parameter removed since was neither
1463  // used nor returned to the user.
1464  return result.full();
1465  }
1466 
1477  const ElementType* query_point, const Size num_closest,
1478  AccessorType* out_indices, DistanceType* out_distances_sq,
1479  const int /* nChecks_IGNORED */ = 10) const
1480  {
1482  num_closest);
1483  resultSet.init(out_indices, out_distances_sq);
1484  this->findNeighbors(resultSet, query_point, nanoflann::SearchParams());
1485  return resultSet.size();
1486  }
1487 
1505  const ElementType* query_point, const DistanceType& radius,
1506  std::vector<std::pair<AccessorType, DistanceType>>& IndicesDists,
1507  const SearchParams& searchParams) const
1508  {
1510  radius, IndicesDists);
1511  const Size nFound =
1512  radiusSearchCustomCallback(query_point, resultSet, searchParams);
1513  if (searchParams.sorted)
1514  std::sort(
1515  IndicesDists.begin(), IndicesDists.end(), IndexDist_Sorter());
1516  return nFound;
1517  }
1518 
1524  template <class SEARCH_CALLBACK>
1526  const ElementType* query_point, SEARCH_CALLBACK& resultSet,
1527  const SearchParams& searchParams = SearchParams()) const
1528  {
1529  this->findNeighbors(resultSet, query_point, searchParams);
1530  return resultSet.size();
1531  }
1532 
1535  public:
1538  void init_vind()
1539  {
1540  // Create a permutable array of indices to the input vectors.
1541  BaseClassRef::m_size = dataset.kdtree_get_point_count();
1542  if (BaseClassRef::vAcc.size() != BaseClassRef::m_size)
1543  BaseClassRef::vAcc.resize(BaseClassRef::m_size);
1544  for (Size i = 0; i < BaseClassRef::m_size; i++)
1545  BaseClassRef::vAcc[i] = i;
1546  }
1547 
1548  void computeBoundingBox(BoundingBox& bbox)
1549  {
1550  resize(bbox, (DIM > 0 ? DIM : BaseClassRef::dim));
1551  if (dataset.kdtree_get_bbox(bbox))
1552  {
1553  // Done! It was implemented in derived class
1554  }
1555  else
1556  {
1557  const Size N = dataset.kdtree_get_point_count();
1558  if (!N)
1559  throw std::runtime_error(
1560  "[nanoflann] computeBoundingBox() called but "
1561  "no data points found.");
1562  for (Dimension i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim); ++i)
1563  {
1564  bbox[i].low = bbox[i].high =
1565  this->dataset_get(*this, BaseClassRef::vAcc[0], i);
1566  }
1567  for (Offset k = 1; k < N; ++k)
1568  {
1569  for (Dimension i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim);
1570  ++i)
1571  {
1572  if (this->dataset_get(*this, BaseClassRef::vAcc[k], i) <
1573  bbox[i].low)
1574  bbox[i].low =
1575  this->dataset_get(*this, BaseClassRef::vAcc[k], i);
1576  if (this->dataset_get(*this, BaseClassRef::vAcc[k], i) >
1577  bbox[i].high)
1578  bbox[i].high =
1579  this->dataset_get(*this, BaseClassRef::vAcc[k], i);
1580  }
1581  }
1582  }
1583  }
1584 
1591  template <class RESULTSET>
1593  RESULTSET& result_set, const ElementType* vec, const NodePtr node,
1594  DistanceType mindistsq, distance_vector_t& dists,
1595  const float epsError) const
1596  {
1597  /* If this is a leaf node, then do check and return. */
1598  if ((node->child1 == nullptr) && (node->child2 == nullptr))
1599  {
1600  // count_leaf += (node->lr.right-node->lr.left); // Removed since
1601  // was neither used nor returned to the user.
1602  DistanceType worst_dist = result_set.worstDist();
1603  for (Offset i = node->node_type.lr.left;
1604  i < node->node_type.lr.right; ++i)
1605  {
1606  const AccessorType accessor =
1607  BaseClassRef::vAcc[i]; // reorder... : i;
1608  DistanceType dist = distance.evalMetric(
1609  vec, accessor, (DIM > 0 ? DIM : BaseClassRef::dim));
1610  if (dist < worst_dist)
1611  {
1612  if (!result_set.addPoint(dist, BaseClassRef::vAcc[i]))
1613  {
1614  // the resultset doesn't want to receive any more
1615  // points, we're done searching!
1616  return false;
1617  }
1618  }
1619  }
1620  return true;
1621  }
1622 
1623  /* Which child branch should be taken first? */
1624  Dimension idx = node->node_type.sub.divfeat;
1625  ElementType val = vec[idx];
1626  DistanceType diff1 = val - node->node_type.sub.divlow;
1627  DistanceType diff2 = val - node->node_type.sub.divhigh;
1628 
1629  NodePtr bestChild;
1630  NodePtr otherChild;
1631  DistanceType cut_dist;
1632  if ((diff1 + diff2) < 0)
1633  {
1634  bestChild = node->child1;
1635  otherChild = node->child2;
1636  cut_dist =
1637  distance.accum_dist(val, node->node_type.sub.divhigh, idx);
1638  }
1639  else
1640  {
1641  bestChild = node->child2;
1642  otherChild = node->child1;
1643  cut_dist =
1644  distance.accum_dist(val, node->node_type.sub.divlow, idx);
1645  }
1646 
1647  /* Call recursively to search next level down. */
1648  if (!searchLevel(
1649  result_set, vec, bestChild, mindistsq, dists, epsError))
1650  {
1651  // the resultset doesn't want to receive any more points, we're done
1652  // searching!
1653  return false;
1654  }
1655 
1656  DistanceType dst = dists[idx];
1657  mindistsq = mindistsq + cut_dist - dst;
1658  dists[idx] = cut_dist;
1659  if (mindistsq * epsError <= result_set.worstDist())
1660  {
1661  if (!searchLevel(
1662  result_set, vec, otherChild, mindistsq, dists, epsError))
1663  {
1664  // the resultset doesn't want to receive any more points, we're
1665  // done searching!
1666  return false;
1667  }
1668  }
1669  dists[idx] = dst;
1670  return true;
1671  }
1672 
1673  public:
1679  void saveIndex(std::ostream& stream) { this->saveIndex_(*this, stream); }
1680 
1686  void loadIndex(std::istream& stream) { this->loadIndex_(*this, stream); }
1687 
1688 }; // class KDTree
1689 
1726 template <
1727  typename Distance, class DatasetAdaptor, int32_t DIM = -1,
1728  typename AccessorType = uint32_t>
1730  : public KDTreeBaseClass<
1731  KDTreeSingleIndexDynamicAdaptor_<
1732  Distance, DatasetAdaptor, DIM, AccessorType>,
1733  Distance, DatasetAdaptor, DIM, AccessorType>
1734 {
1735  public:
1739  const DatasetAdaptor& dataset;
1740 
1741  KDTreeSingleIndexAdaptorParams index_params;
1742 
1743  std::vector<int>& treeIndex;
1744 
1745  Distance distance;
1746 
1747  using BaseClassRef = typename nanoflann::KDTreeBaseClass<
1749  Distance, DatasetAdaptor, DIM, AccessorType>,
1750  Distance, DatasetAdaptor, DIM, AccessorType>;
1751 
1752  using ElementType = typename BaseClassRef::ElementType;
1753  using DistanceType = typename BaseClassRef::DistanceType;
1754 
1755  using Offset = typename BaseClassRef::Offset;
1756  using Size = typename BaseClassRef::Size;
1757  using Dimension = typename BaseClassRef::Dimension;
1758 
1759  using Node = typename BaseClassRef::Node;
1760  using NodePtr = Node*;
1761 
1762  using Interval = typename BaseClassRef::Interval;
1765  using BoundingBox = typename BaseClassRef::BoundingBox;
1766 
1769  using distance_vector_t = typename BaseClassRef::distance_vector_t;
1770 
1787  const Dimension dimensionality, const DatasetAdaptor& inputData,
1788  std::vector<int>& treeIndex_,
1789  const KDTreeSingleIndexAdaptorParams& params =
1791  : dataset(inputData),
1792  index_params(params),
1793  treeIndex(treeIndex_),
1794  distance(inputData)
1795  {
1796  BaseClassRef::root_node = nullptr;
1797  BaseClassRef::m_size = 0;
1798  BaseClassRef::m_size_at_index_build = 0;
1799  BaseClassRef::dim = dimensionality;
1800  if (DIM > 0) BaseClassRef::dim = DIM;
1801  BaseClassRef::m_leaf_max_size = params.leaf_max_size;
1802  }
1803 
1806  const KDTreeSingleIndexDynamicAdaptor_& rhs) = default;
1807 
1811  {
1813  std::swap(BaseClassRef::vAcc, tmp.BaseClassRef::vAcc);
1814  std::swap(
1815  BaseClassRef::m_leaf_max_size, tmp.BaseClassRef::m_leaf_max_size);
1816  std::swap(index_params, tmp.index_params);
1817  std::swap(treeIndex, tmp.treeIndex);
1818  std::swap(BaseClassRef::m_size, tmp.BaseClassRef::m_size);
1819  std::swap(
1820  BaseClassRef::m_size_at_index_build,
1821  tmp.BaseClassRef::m_size_at_index_build);
1822  std::swap(BaseClassRef::root_node, tmp.BaseClassRef::root_node);
1823  std::swap(BaseClassRef::root_bbox, tmp.BaseClassRef::root_bbox);
1824  std::swap(BaseClassRef::pool, tmp.BaseClassRef::pool);
1825  return *this;
1826  }
1827 
1831  void buildIndex()
1832  {
1833  BaseClassRef::m_size = BaseClassRef::vAcc.size();
1834  this->freeIndex(*this);
1835  BaseClassRef::m_size_at_index_build = BaseClassRef::m_size;
1836  if (BaseClassRef::m_size == 0) return;
1837  computeBoundingBox(BaseClassRef::root_bbox);
1838  BaseClassRef::root_node = this->divideTree(
1839  *this, 0, BaseClassRef::m_size,
1840  BaseClassRef::root_bbox); // construct the tree
1841  }
1842 
1859  template <typename RESULTSET>
1861  RESULTSET& result, const ElementType* vec,
1862  const SearchParams& searchParams) const
1863  {
1864  assert(vec);
1865  if (this->size(*this) == 0) return false;
1866  if (!BaseClassRef::root_node) return false;
1867  float epsError = 1 + searchParams.eps;
1868 
1869  // fixed or variable-sized container (depending on DIM)
1870  distance_vector_t dists;
1871  // Fill it with zeros.
1872  assign(
1873  dists, (DIM > 0 ? DIM : BaseClassRef::dim),
1874  static_cast<typename distance_vector_t::value_type>(0));
1875  DistanceType distsq = this->computeInitialDistances(*this, vec, dists);
1876  searchLevel(
1877  result, vec, BaseClassRef::root_node, distsq, dists,
1878  epsError); // "count_leaf" parameter removed since was neither
1879  // used nor returned to the user.
1880  return result.full();
1881  }
1882 
1893  const ElementType* query_point, const Size num_closest,
1894  AccessorType* out_indices, DistanceType* out_distances_sq,
1895  const int /* nChecks_IGNORED */ = 10) const
1896  {
1898  num_closest);
1899  resultSet.init(out_indices, out_distances_sq);
1900  this->findNeighbors(resultSet, query_point, nanoflann::SearchParams());
1901  return resultSet.size();
1902  }
1903 
1921  const ElementType* query_point, const DistanceType& radius,
1922  std::vector<std::pair<AccessorType, DistanceType>>& IndicesDists,
1923  const SearchParams& searchParams) const
1924  {
1926  radius, IndicesDists);
1927  const size_t nFound =
1928  radiusSearchCustomCallback(query_point, resultSet, searchParams);
1929  if (searchParams.sorted)
1930  std::sort(
1931  IndicesDists.begin(), IndicesDists.end(), IndexDist_Sorter());
1932  return nFound;
1933  }
1934 
1940  template <class SEARCH_CALLBACK>
1942  const ElementType* query_point, SEARCH_CALLBACK& resultSet,
1943  const SearchParams& searchParams = SearchParams()) const
1944  {
1945  this->findNeighbors(resultSet, query_point, searchParams);
1946  return resultSet.size();
1947  }
1948 
1951  public:
1952  void computeBoundingBox(BoundingBox& bbox)
1953  {
1954  resize(bbox, (DIM > 0 ? DIM : BaseClassRef::dim));
1955 
1956  if (dataset.kdtree_get_bbox(bbox))
1957  {
1958  // Done! It was implemented in derived class
1959  }
1960  else
1961  {
1962  const Size N = BaseClassRef::m_size;
1963  if (!N)
1964  throw std::runtime_error(
1965  "[nanoflann] computeBoundingBox() called but "
1966  "no data points found.");
1967  for (Dimension i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim); ++i)
1968  {
1969  bbox[i].low = bbox[i].high =
1970  this->dataset_get(*this, BaseClassRef::vAcc[0], i);
1971  }
1972  for (Offset k = 1; k < N; ++k)
1973  {
1974  for (Dimension i = 0; i < (DIM > 0 ? DIM : BaseClassRef::dim);
1975  ++i)
1976  {
1977  if (this->dataset_get(*this, BaseClassRef::vAcc[k], i) <
1978  bbox[i].low)
1979  bbox[i].low =
1980  this->dataset_get(*this, BaseClassRef::vAcc[k], i);
1981  if (this->dataset_get(*this, BaseClassRef::vAcc[k], i) >
1982  bbox[i].high)
1983  bbox[i].high =
1984  this->dataset_get(*this, BaseClassRef::vAcc[k], i);
1985  }
1986  }
1987  }
1988  }
1989 
1994  template <class RESULTSET>
1996  RESULTSET& result_set, const ElementType* vec, const NodePtr node,
1997  DistanceType mindistsq, distance_vector_t& dists,
1998  const float epsError) const
1999  {
2000  /* If this is a leaf node, then do check and return. */
2001  if ((node->child1 == nullptr) && (node->child2 == nullptr))
2002  {
2003  // count_leaf += (node->lr.right-node->lr.left); // Removed since
2004  // was neither used nor returned to the user.
2005  DistanceType worst_dist = result_set.worstDist();
2006  for (Offset i = node->node_type.lr.left;
2007  i < node->node_type.lr.right; ++i)
2008  {
2009  const AccessorType index =
2010  BaseClassRef::vAcc[i]; // reorder... : i;
2011  if (treeIndex[index] == -1) continue;
2012  DistanceType dist = distance.evalMetric(
2013  vec, index, (DIM > 0 ? DIM : BaseClassRef::dim));
2014  if (dist < worst_dist)
2015  {
2016  if (!result_set.addPoint(
2017  static_cast<typename RESULTSET::DistanceType>(dist),
2018  static_cast<typename RESULTSET::IndexType>(
2019  BaseClassRef::vAcc[i])))
2020  {
2021  // the resultset doesn't want to receive any more
2022  // points, we're done searching!
2023  return; // false;
2024  }
2025  }
2026  }
2027  return;
2028  }
2029 
2030  /* Which child branch should be taken first? */
2031  Dimension idx = node->node_type.sub.divfeat;
2032  ElementType val = vec[idx];
2033  DistanceType diff1 = val - node->node_type.sub.divlow;
2034  DistanceType diff2 = val - node->node_type.sub.divhigh;
2035 
2036  NodePtr bestChild;
2037  NodePtr otherChild;
2038  DistanceType cut_dist;
2039  if ((diff1 + diff2) < 0)
2040  {
2041  bestChild = node->child1;
2042  otherChild = node->child2;
2043  cut_dist =
2044  distance.accum_dist(val, node->node_type.sub.divhigh, idx);
2045  }
2046  else
2047  {
2048  bestChild = node->child2;
2049  otherChild = node->child1;
2050  cut_dist =
2051  distance.accum_dist(val, node->node_type.sub.divlow, idx);
2052  }
2053 
2054  /* Call recursively to search next level down. */
2055  searchLevel(result_set, vec, bestChild, mindistsq, dists, epsError);
2056 
2057  DistanceType dst = dists[idx];
2058  mindistsq = mindistsq + cut_dist - dst;
2059  dists[idx] = cut_dist;
2060  if (mindistsq * epsError <= result_set.worstDist())
2061  {
2062  searchLevel(
2063  result_set, vec, otherChild, mindistsq, dists, epsError);
2064  }
2065  dists[idx] = dst;
2066  }
2067 
2068  public:
2074  void saveIndex(std::ostream& stream) { this->saveIndex_(*this, stream); }
2075 
2081  void loadIndex(std::istream& stream) { this->loadIndex_(*this, stream); }
2082 };
2083 
2098 template <
2099  typename Distance, class DatasetAdaptor, int32_t DIM = -1,
2100  typename AccessorType = uint32_t>
2102 {
2103  public:
2104  using ElementType = typename Distance::ElementType;
2105  using DistanceType = typename Distance::DistanceType;
2106 
2107  using Offset = typename KDTreeSingleIndexDynamicAdaptor_<
2108  Distance, DatasetAdaptor, DIM>::Offset;
2109  using Size = typename KDTreeSingleIndexDynamicAdaptor_<
2110  Distance, DatasetAdaptor, DIM>::Size;
2111  using Dimension = typename KDTreeSingleIndexDynamicAdaptor_<
2112  Distance, DatasetAdaptor, DIM>::Dimension;
2113 
2114  protected:
2115  Size m_leaf_max_size;
2116  Size treeCount;
2117  Size pointCount;
2118 
2122  const DatasetAdaptor& dataset;
2123 
2124  std::vector<int>
2128  std::unordered_set<int> removedPoints;
2129 
2130  KDTreeSingleIndexAdaptorParams index_params;
2131 
2132  Dimension dim;
2133 
2135  Distance, DatasetAdaptor, DIM, AccessorType>;
2136  std::vector<index_container_t> index;
2137 
2138  public:
2141  const std::vector<index_container_t>& getAllIndices() const
2142  {
2143  return index;
2144  }
2145 
2146  private:
2148  int First0Bit(AccessorType num)
2149  {
2150  int pos = 0;
2151  while (num & 1)
2152  {
2153  num = num >> 1;
2154  pos++;
2155  }
2156  return pos;
2157  }
2158 
2160  void init()
2161  {
2162  using my_kd_tree_t = KDTreeSingleIndexDynamicAdaptor_<
2163  Distance, DatasetAdaptor, DIM, AccessorType>;
2164  std::vector<my_kd_tree_t> index_(
2165  treeCount,
2166  my_kd_tree_t(dim /*dim*/, dataset, treeIndex, index_params));
2167  index = index_;
2168  }
2169 
2170  public:
2171  Distance distance;
2172 
2189  const int dimensionality, const DatasetAdaptor& inputData,
2190  const KDTreeSingleIndexAdaptorParams& params =
2192  const size_t maximumPointCount = 1000000000U)
2193  : dataset(inputData), index_params(params), distance(inputData)
2194  {
2195  treeCount = static_cast<size_t>(std::log2(maximumPointCount)) + 1;
2196  pointCount = 0U;
2197  dim = dimensionality;
2198  treeIndex.clear();
2199  if (DIM > 0) dim = DIM;
2200  m_leaf_max_size = params.leaf_max_size;
2201  init();
2202  const size_t num_initial_points = dataset.kdtree_get_point_count();
2203  if (num_initial_points > 0) { addPoints(0, num_initial_points - 1); }
2204  }
2205 
2209  Distance, DatasetAdaptor, DIM, AccessorType>&) = delete;
2210 
2212  void addPoints(AccessorType start, AccessorType end)
2213  {
2214  const Size count = end - start + 1;
2215  int maxIndex = 0;
2216  treeIndex.resize(treeIndex.size() + count);
2217  for (AccessorType idx = start; idx <= end; idx++)
2218  {
2219  const int pos = First0Bit(pointCount);
2220  maxIndex = std::max(pos, maxIndex);
2221  treeIndex[pointCount] = pos;
2222 
2223  const auto it = removedPoints.find(idx);
2224  if (it != removedPoints.end())
2225  {
2226  removedPoints.erase(it);
2227  treeIndex[idx] = pos;
2228  }
2229 
2230  for (int i = 0; i < pos; i++)
2231  {
2232  for (int j = 0; j < static_cast<int>(index[i].vAcc.size()); j++)
2233  {
2234  index[pos].vAcc.push_back(index[i].vAcc[j]);
2235  if (treeIndex[index[i].vAcc[j]] != -1)
2236  treeIndex[index[i].vAcc[j]] = pos;
2237  }
2238  index[i].vAcc.clear();
2239  }
2240  index[pos].vAcc.push_back(idx);
2241  pointCount++;
2242  }
2243 
2244  for (int i = 0; i <= maxIndex; ++i)
2245  {
2246  index[i].freeIndex(index[i]);
2247  if (!index[i].vAcc.empty()) index[i].buildIndex();
2248  }
2249  }
2250 
2252  void removePoint(size_t idx)
2253  {
2254  if (idx >= pointCount) return;
2255  removedPoints.insert(idx);
2256  treeIndex[idx] = -1;
2257  }
2258 
2272  template <typename RESULTSET>
2274  RESULTSET& result, const ElementType* vec,
2275  const SearchParams& searchParams) const
2276  {
2277  for (size_t i = 0; i < treeCount; i++)
2278  {
2279  index[i].findNeighbors(result, &vec[0], searchParams);
2280  }
2281  return result.full();
2282  }
2283 };
2284 
2309 template <
2310  class MatrixType, int32_t DIM = -1, class Distance = nanoflann::metric_L2,
2311  bool row_major = true>
2313 {
2314  using self_t =
2316  using num_t = typename MatrixType::Scalar;
2317  using IndexType = typename MatrixType::Index;
2318  using metric_t = typename Distance::template traits<
2319  num_t, self_t, IndexType>::distance_t;
2320 
2322  metric_t, self_t,
2323  row_major ? MatrixType::ColsAtCompileTime
2324  : MatrixType::RowsAtCompileTime,
2325  IndexType>;
2326 
2327  index_t* index;
2329 
2330  using Offset = typename index_t::Offset;
2331  using Size = typename index_t::Size;
2332  using Dimension = typename index_t::Dimension;
2333 
2336  const Dimension dimensionality,
2337  const std::reference_wrapper<const MatrixType>& mat,
2338  const int leaf_max_size = 10)
2339  : m_data_matrix(mat)
2340  {
2341  const auto dims = row_major ? mat.get().cols() : mat.get().rows();
2342  if (static_cast<Dimension>(dims) != dimensionality)
2343  throw std::runtime_error(
2344  "Error: 'dimensionality' must match column count in data "
2345  "matrix");
2346  if (DIM > 0 && static_cast<int32_t>(dims) != DIM)
2347  throw std::runtime_error(
2348  "Data set dimensionality does not match the 'DIM' template "
2349  "argument");
2350  index = new index_t(
2351  dims, *this /* adaptor */,
2353  }
2354 
2355  public:
2358 
2359  ~KDTreeEigenMatrixAdaptor() { delete index; }
2360 
2361  const std::reference_wrapper<const MatrixType> m_data_matrix;
2362 
2369  inline void query(
2370  const num_t* query_point, const Size num_closest,
2371  IndexType* out_indices, num_t* out_distances_sq,
2372  const int /* nChecks_IGNORED */ = 10) const
2373  {
2374  nanoflann::KNNResultSet<num_t, IndexType> resultSet(num_closest);
2375  resultSet.init(out_indices, out_distances_sq);
2376  index->findNeighbors(resultSet, query_point, nanoflann::SearchParams());
2377  }
2378 
2382  const self_t& derived() const { return *this; }
2383  self_t& derived() { return *this; }
2384 
2385  // Must return the number of data points
2386  inline Size kdtree_get_point_count() const
2387  {
2388  if (row_major)
2389  return m_data_matrix.get().rows();
2390  else
2391  return m_data_matrix.get().cols();
2392  }
2393 
2394  // Returns the dim'th component of the idx'th point in the class:
2395  inline num_t kdtree_get_pt(const IndexType idx, size_t dim) const
2396  {
2397  if (row_major)
2398  return m_data_matrix.get().coeff(idx, IndexType(dim));
2399  else
2400  return m_data_matrix.get().coeff(IndexType(dim), idx);
2401  }
2402 
2403  // Optional bounding-box computation: return false to default to a standard
2404  // bbox computation loop.
2405  // Return true if the BBOX was already computed by the class and returned
2406  // in "bb" so it can be avoided to redo it again. Look at bb.size() to
2407  // find out the expected dimensionality (e.g. 2 or 3 for point clouds)
2408  template <class BBOX>
2409  bool kdtree_get_bbox(BBOX& /*bb*/) const
2410  {
2411  return false;
2412  }
2413 
2416 }; // end of KDTreeEigenMatrixAdaptor // end of grouping
2420 } // namespace nanoflann
Definition: nanoflann.hpp:896
Size veclen(const Derived &obj)
Definition: nanoflann.hpp:985
NodePtr divideTree(Derived &obj, const Offset left, const Offset right, BoundingBox &bbox)
Definition: nanoflann.hpp:1026
PooledAllocator pool
Definition: nanoflann.hpp:979
Dimension dim
Dimensionality of each data point.
Definition: nanoflann.hpp:956
Size m_size_at_index_build
Definition: nanoflann.hpp:954
ElementType dataset_get(const Derived &obj, AccessorType element, Dimension component) const
Helper accessor to the dataset points:
Definition: nanoflann.hpp:988
Size size(const Derived &obj) const
Definition: nanoflann.hpp:982
BoundingBox root_bbox
Definition: nanoflann.hpp:970
Size usedMemory(Derived &obj)
Definition: nanoflann.hpp:998
void saveIndex_(Derived &obj, std::ostream &stream)
Definition: nanoflann.hpp:1233
std::vector< AccessorType > vAcc
Definition: nanoflann.hpp:913
void planeSplit(Derived &obj, Offset ind, const Size count, Dimension cutfeat, DistanceType &cutval, Offset &lim1, Offset &lim2)
Definition: nanoflann.hpp:1145
void freeIndex(Derived &obj)
Definition: nanoflann.hpp:900
void loadIndex_(Derived &obj, std::istream &stream)
Definition: nanoflann.hpp:1248
typename array_or_vector_selector< DIM, Interval >::container_t BoundingBox
Definition: nanoflann.hpp:961
Size m_size
Number of current points in the dataset.
Definition: nanoflann.hpp:953
typename array_or_vector_selector< DIM, DistanceType >::container_t distance_vector_t
Definition: nanoflann.hpp:966
Definition: nanoflann.hpp:1307
KDTreeSingleIndexAdaptor(const Dimension dimensionality, const DatasetAdaptor &inputData, const KDTreeSingleIndexAdaptorParams &params, Args &&... args)
Definition: nanoflann.hpp:1369
void saveIndex(std::ostream &stream)
Definition: nanoflann.hpp:1679
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< std::pair< AccessorType, DistanceType >> &IndicesDists, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1504
void buildIndex()
Definition: nanoflann.hpp:1410
bool searchLevel(RESULTSET &result_set, const ElementType *vec, const NodePtr node, DistanceType mindistsq, distance_vector_t &dists, const float epsError) const
Definition: nanoflann.hpp:1592
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:1525
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1441
typename BaseClassRef::distance_vector_t distance_vector_t
Definition: nanoflann.hpp:1346
Size knnSearch(const ElementType *query_point, const Size num_closest, AccessorType *out_indices, DistanceType *out_distances_sq, const int=10) const
Definition: nanoflann.hpp:1476
void init_vind()
Definition: nanoflann.hpp:1538
void loadIndex(std::istream &stream)
Definition: nanoflann.hpp:1686
const DatasetAdaptor & dataset
The source of our data.
Definition: nanoflann.hpp:1317
typename BaseClassRef::BoundingBox BoundingBox
Definition: nanoflann.hpp:1342
Definition: nanoflann.hpp:1734
const DatasetAdaptor & dataset
The source of our data.
Definition: nanoflann.hpp:1739
void searchLevel(RESULTSET &result_set, const ElementType *vec, const NodePtr node, DistanceType mindistsq, distance_vector_t &dists, const float epsError) const
Definition: nanoflann.hpp:1995
Size radiusSearchCustomCallback(const ElementType *query_point, SEARCH_CALLBACK &resultSet, const SearchParams &searchParams=SearchParams()) const
Definition: nanoflann.hpp:1941
typename BaseClassRef::distance_vector_t distance_vector_t
Definition: nanoflann.hpp:1769
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< std::pair< AccessorType, DistanceType >> &IndicesDists, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1920
void buildIndex()
Definition: nanoflann.hpp:1831
void loadIndex(std::istream &stream)
Definition: nanoflann.hpp:2081
KDTreeSingleIndexDynamicAdaptor_(const KDTreeSingleIndexDynamicAdaptor_ &rhs)=default
typename BaseClassRef::BoundingBox BoundingBox
Definition: nanoflann.hpp:1765
KDTreeSingleIndexDynamicAdaptor_ operator=(const KDTreeSingleIndexDynamicAdaptor_ &rhs)
Definition: nanoflann.hpp:1809
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1860
KDTreeSingleIndexDynamicAdaptor_(const Dimension dimensionality, const DatasetAdaptor &inputData, std::vector< int > &treeIndex_, const KDTreeSingleIndexAdaptorParams &params=KDTreeSingleIndexAdaptorParams())
Definition: nanoflann.hpp:1786
Size knnSearch(const ElementType *query_point, const Size num_closest, AccessorType *out_indices, DistanceType *out_distances_sq, const int=10) const
Definition: nanoflann.hpp:1892
void saveIndex(std::ostream &stream)
Definition: nanoflann.hpp:2074
Definition: nanoflann.hpp:2102
void removePoint(size_t idx)
Definition: nanoflann.hpp:2252
const DatasetAdaptor & dataset
The source of our data.
Definition: nanoflann.hpp:2122
std::vector< int > treeIndex
Definition: nanoflann.hpp:2125
void addPoints(AccessorType start, AccessorType end)
Definition: nanoflann.hpp:2212
Dimension dim
Dimensionality of each data point.
Definition: nanoflann.hpp:2132
KDTreeSingleIndexDynamicAdaptor(const KDTreeSingleIndexDynamicAdaptor< Distance, DatasetAdaptor, DIM, AccessorType > &)=delete
const std::vector< index_container_t > & getAllIndices() const
Definition: nanoflann.hpp:2141
KDTreeSingleIndexDynamicAdaptor(const int dimensionality, const DatasetAdaptor &inputData, const KDTreeSingleIndexAdaptorParams &params=KDTreeSingleIndexAdaptorParams(), const size_t maximumPointCount=1000000000U)
Definition: nanoflann.hpp:2188
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParams &searchParams) const
Definition: nanoflann.hpp:2273
Definition: nanoflann.hpp:163
bool addPoint(DistanceType dist, IndexType index)
Definition: nanoflann.hpp:199
Definition: nanoflann.hpp:740
T * allocate(const size_t count=1)
Definition: nanoflann.hpp:851
~PooledAllocator()
Definition: nanoflann.hpp:775
void free_all()
Definition: nanoflann.hpp:778
PooledAllocator()
Definition: nanoflann.hpp:770
void * malloc(const size_t req_size)
Definition: nanoflann.hpp:794
Definition: nanoflann.hpp:253
std::pair< IndexType, DistanceType > worst_item() const
Definition: nanoflann.hpp:296
bool addPoint(DistanceType dist, IndexType index)
Definition: nanoflann.hpp:283
const size_t WORDSIZE
Definition: nanoflann.hpp:736
T * allocate(size_t count=1)
Definition: nanoflann.hpp:715
T pi_const()
Definition: nanoflann.hpp:84
std::enable_if< has_assign< Container >::value, void >::type assign(Container &c, const size_t nElements, const T &value)
Definition: nanoflann.hpp:141
std::enable_if< has_resize< Container >::value, void >::type resize(Container &c, const size_t nElements)
Definition: nanoflann.hpp:119
Definition: nanoflann.hpp:239
bool operator()(const PairType &p1, const PairType &p2) const
Definition: nanoflann.hpp:242
Definition: nanoflann.hpp:945
Definition: nanoflann.hpp:922
DistanceType divhigh
The values used for subdivision.
Definition: nanoflann.hpp:935
Dimension divfeat
Dimension used for subdivision.
Definition: nanoflann.hpp:933
Offset right
Indices of points in leaf node.
Definition: nanoflann.hpp:929
union nanoflann::KDTreeBaseClass::Node::@0 node_type
Node * child1
Definition: nanoflann.hpp:939
Definition: nanoflann.hpp:2313
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:2369
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:2335
KDTreeEigenMatrixAdaptor(const self_t &)=delete
typename index_t::Offset Offset
Definition: nanoflann.hpp:2330
Definition: nanoflann.hpp:674
Definition: nanoflann.hpp:365
Definition: nanoflann.hpp:427
Definition: nanoflann.hpp:492
Definition: nanoflann.hpp:348
Definition: nanoflann.hpp:537
DistanceType accum_dist(const U a, const V b, const size_t) const
Definition: nanoflann.hpp:555
Definition: nanoflann.hpp:582
Definition: nanoflann.hpp:688
bool sorted
distance (default: true)
Definition: nanoflann.hpp:699
SearchParams(int checks_IGNORED_=32, float eps_=0, bool sorted_=true)
Definition: nanoflann.hpp:691
float eps
search for eps-approximate neighbours (default: 0)
Definition: nanoflann.hpp:698
int checks
Definition: nanoflann.hpp:696
Definition: nanoflann.hpp:867
Definition: nanoflann.hpp:106
Definition: nanoflann.hpp:95
Definition: nanoflann.hpp:612
Definition: nanoflann.hpp:609
Definition: nanoflann.hpp:621
Definition: nanoflann.hpp:630
Definition: nanoflann.hpp:627
Definition: nanoflann.hpp:618
Definition: nanoflann.hpp:639
Definition: nanoflann.hpp:636
Definition: nanoflann.hpp:648
Definition: nanoflann.hpp:645