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
73namespace nanoflann
74{
79template <typename T>
81{
82 return static_cast<T>(3.14159265358979323846);
83}
84
89template <typename T, typename = int>
90struct has_resize : std::false_type
91{
92};
93
94template <typename T>
95struct has_resize<T, decltype((void)std::declval<T>().resize(1), 0)>
96 : std::true_type
97{
98};
99
100template <typename T, typename = int>
101struct has_assign : std::false_type
102{
103};
104
105template <typename T>
106struct has_assign<T, decltype((void)std::declval<T>().assign(1, 0), 0)>
107 : std::true_type
108{
109};
110
114template <typename Container>
115inline 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
125template <typename Container>
126inline 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
136template <typename Container, typename T>
137inline 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
146template <typename Container, typename T>
147inline 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
155template <
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
247template <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
310template <typename T>
311void save_value(std::ostream& stream, const T& value)
312{
313 stream.write(reinterpret_cast<const char*>(&value), sizeof(T));
314}
315
316template <typename T>
317void 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
324template <typename T>
325void load_value(std::istream& stream, T& value)
326{
327 stream.read(reinterpret_cast<char*>(&value), sizeof(T));
328}
329
330template <typename T>
331void 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}
343struct Metric
344{
345};
346
357template <
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
419template <
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
484template <
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
529template <
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
574template <
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
604struct metric_L1 : public Metric
605{
606 template <class T, class DataSource, typename AccessorType = uint32_t>
607 struct traits
608 {
610 };
611};
613struct metric_L2 : public Metric
614{
615 template <class T, class DataSource, typename AccessorType = uint32_t>
616 struct traits
617 {
619 };
620};
623{
624 template <class T, class DataSource, typename AccessorType = uint32_t>
625 struct traits
626 {
628 };
629};
631struct metric_SO2 : public Metric
632{
633 template <class T, class DataSource, typename AccessorType = uint32_t>
634 struct traits
635 {
637 };
638};
640struct metric_SO3 : public Metric
641{
642 template <class T, class DataSource, typename AccessorType = uint32_t>
643 struct traits
644 {
646 };
647};
648
654enum class KDTreeSingleIndexAdaptorFlags
655{
656 None = 0,
657 SkipInitialBuildIndex = 1
658};
659
660inline 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};
710template <typename T>
711inline T* allocate(size_t count = 1)
712{
713 T* mem = static_cast<T*>(::malloc(sizeof(T) * count));
714 return mem;
715}
716
732const size_t WORDSIZE = 16;
733const 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};
861template <int32_t DIM, typename T>
863{
864 using container_t = std::array<T, DIM>;
865};
867template <typename T>
869{
870 using container_t = std::vector<T>;
871};
872
888template <
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
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
1296template <
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:
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:
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
1722template <
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
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
2094template <
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
2305template <
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
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< std::pair< AccessorType, DistanceType > > &IndicesDists, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1500
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
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
void buildIndex()
Definition: nanoflann.hpp:1827
void loadIndex(std::istream &stream)
Definition: nanoflann.hpp:2077
KDTreeSingleIndexDynamicAdaptor_(const KDTreeSingleIndexDynamicAdaptor_ &rhs)=default
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< std::pair< AccessorType, DistanceType > > &IndicesDists, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1916
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
~PooledAllocator()
Definition: nanoflann.hpp:771
void free_all()
Definition: nanoflann.hpp:774
void * malloc(const size_t req_size)
Definition: nanoflann.hpp:790
T * allocate(const size_t count=1)
Definition: nanoflann.hpp:847
PooledAllocator()
Definition: nanoflann.hpp:766
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
std::enable_if< has_assign< Container >::value, void >::type assign(Container &c, const size_t nElements, const T &value)
Definition: nanoflann.hpp:137
T pi_const()
Definition: nanoflann.hpp:80
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