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
77namespace nanoflann
78{
83template <typename T>
85{
86 return static_cast<T>(3.14159265358979323846);
87}
88
93template <typename T, typename = int>
94struct has_resize : std::false_type
95{
96};
97
98template <typename T>
99struct has_resize<T, decltype((void)std::declval<T>().resize(1), 0)>
100 : std::true_type
101{
102};
103
104template <typename T, typename = int>
105struct has_assign : std::false_type
106{
107};
108
109template <typename T>
110struct has_assign<T, decltype((void)std::declval<T>().assign(1, 0), 0)>
111 : std::true_type
112{
113};
114
118template <typename Container>
119inline 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
129template <typename Container>
130inline 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
140template <typename Container, typename T>
141inline 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
150template <typename Container, typename T>
151inline 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
159template <
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
251template <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
314template <typename T>
315void save_value(std::ostream& stream, const T& value)
316{
317 stream.write(reinterpret_cast<const char*>(&value), sizeof(T));
318}
319
320template <typename T>
321void 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
328template <typename T>
329void load_value(std::istream& stream, T& value)
330{
331 stream.read(reinterpret_cast<char*>(&value), sizeof(T));
332}
333
334template <typename T>
335void 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}
347struct Metric
348{
349};
350
361template <
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
423template <
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
488template <
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
533template <
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
578template <
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
608struct metric_L1 : public Metric
609{
610 template <class T, class DataSource, typename AccessorType = uint32_t>
611 struct traits
612 {
614 };
615};
617struct metric_L2 : public Metric
618{
619 template <class T, class DataSource, typename AccessorType = uint32_t>
620 struct traits
621 {
623 };
624};
627{
628 template <class T, class DataSource, typename AccessorType = uint32_t>
629 struct traits
630 {
632 };
633};
635struct metric_SO2 : public Metric
636{
637 template <class T, class DataSource, typename AccessorType = uint32_t>
638 struct traits
639 {
641 };
642};
644struct metric_SO3 : public Metric
645{
646 template <class T, class DataSource, typename AccessorType = uint32_t>
647 struct traits
648 {
650 };
651};
652
658enum class KDTreeSingleIndexAdaptorFlags
659{
660 None = 0,
661 SkipInitialBuildIndex = 1
662};
663
664inline 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};
714template <typename T>
715inline T* allocate(size_t count = 1)
716{
717 T* mem = static_cast<T*>(::malloc(sizeof(T) * count));
718 return mem;
719}
720
736const size_t WORDSIZE = 16;
737const 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};
865template <int32_t DIM, typename T>
867{
868 using container_t = std::array<T, DIM>;
869};
871template <typename T>
873{
874 using container_t = std::vector<T>;
875};
876
892template <
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
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
1300template <
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:
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:
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
1726template <
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
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
2098template <
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
2309template <
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
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< std::pair< AccessorType, DistanceType > > &IndicesDists, const SearchParams &searchParams) const
Definition: nanoflann.hpp:1504
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
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
void buildIndex()
Definition: nanoflann.hpp:1831
void loadIndex(std::istream &stream)
Definition: nanoflann.hpp:2081
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:1920
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
~PooledAllocator()
Definition: nanoflann.hpp:775
void free_all()
Definition: nanoflann.hpp:778
void * malloc(const size_t req_size)
Definition: nanoflann.hpp:794
T * allocate(const size_t count=1)
Definition: nanoflann.hpp:851
PooledAllocator()
Definition: nanoflann.hpp:770
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
std::enable_if< has_assign< Container >::value, void >::type assign(Container &c, const size_t nElements, const T &value)
Definition: nanoflann.hpp:141
T pi_const()
Definition: nanoflann.hpp:84
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