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 <atomic>
50#include <cassert>
51#include <cmath> // for abs()
52#include <cstdlib> // for abs()
53#include <functional> // std::reference_wrapper
54#include <future>
55#include <istream>
56#include <limits> // std::numeric_limits
57#include <ostream>
58#include <stdexcept>
59#include <unordered_set>
60#include <vector>
61
63#define NANOFLANN_VERSION 0x150
64
65// Avoid conflicting declaration of min/max macros in Windows headers
66#if !defined(NOMINMAX) && \
67 (defined(_WIN32) || defined(_WIN32_) || defined(WIN32) || defined(_WIN64))
68#define NOMINMAX
69#ifdef max
70#undef max
71#undef min
72#endif
73#endif
74// Avoid conflicts with X11 headers
75#ifdef None
76#undef None
77#endif
78
79namespace nanoflann
80{
85template <typename T>
87{
88 return static_cast<T>(3.14159265358979323846);
89}
90
95template <typename T, typename = int>
96struct has_resize : std::false_type
97{
98};
99
100template <typename T>
101struct has_resize<T, decltype((void)std::declval<T>().resize(1), 0)>
102 : std::true_type
103{
104};
105
106template <typename T, typename = int>
107struct has_assign : std::false_type
108{
109};
110
111template <typename T>
112struct has_assign<T, decltype((void)std::declval<T>().assign(1, 0), 0)>
113 : std::true_type
114{
115};
116
120template <typename Container>
121inline typename std::enable_if<has_resize<Container>::value, void>::type resize(
122 Container& c, const size_t nElements)
123{
124 c.resize(nElements);
125}
126
131template <typename Container>
132inline typename std::enable_if<!has_resize<Container>::value, void>::type
133 resize(Container& c, const size_t nElements)
134{
135 if (nElements != c.size())
136 throw std::logic_error("Try to change the size of a std::array.");
137}
138
142template <typename Container, typename T>
143inline typename std::enable_if<has_assign<Container>::value, void>::type assign(
144 Container& c, const size_t nElements, const T& value)
145{
146 c.assign(nElements, value);
147}
148
152template <typename Container, typename T>
153inline typename std::enable_if<!has_assign<Container>::value, void>::type
154 assign(Container& c, const size_t nElements, const T& value)
155{
156 for (size_t i = 0; i < nElements; i++) c[i] = value;
157}
158
161template <
162 typename _DistanceType, typename _IndexType = size_t,
163 typename _CountType = size_t>
165{
166 public:
167 using DistanceType = _DistanceType;
168 using IndexType = _IndexType;
169 using CountType = _CountType;
170
171 private:
172 IndexType* indices;
173 DistanceType* dists;
174 CountType capacity;
175 CountType count;
176
177 public:
178 explicit KNNResultSet(CountType capacity_)
179 : indices(nullptr), dists(nullptr), capacity(capacity_), count(0)
180 {
181 }
182
183 void init(IndexType* indices_, DistanceType* dists_)
184 {
185 indices = indices_;
186 dists = dists_;
187 count = 0;
188 if (capacity)
189 dists[capacity - 1] = (std::numeric_limits<DistanceType>::max)();
190 }
191
192 CountType size() const { return count; }
193
194 bool full() const { return count == capacity; }
195
201 bool addPoint(DistanceType dist, IndexType index)
202 {
203 CountType i;
204 for (i = count; i > 0; --i)
205 {
208#ifdef NANOFLANN_FIRST_MATCH
209 if ((dists[i - 1] > dist) ||
210 ((dist == dists[i - 1]) && (indices[i - 1] > index)))
211 {
212#else
213 if (dists[i - 1] > dist)
214 {
215#endif
216 if (i < capacity)
217 {
218 dists[i] = dists[i - 1];
219 indices[i] = indices[i - 1];
220 }
221 }
222 else
223 break;
224 }
225 if (i < capacity)
226 {
227 dists[i] = dist;
228 indices[i] = index;
229 }
230 if (count < capacity) count++;
231
232 // tell caller that the search shall continue
233 return true;
234 }
235
236 DistanceType worstDist() const { return dists[capacity - 1]; }
237};
238
241{
243 template <typename PairType>
244 bool operator()(const PairType& p1, const PairType& p2) const
245 {
246 return p1.second < p2.second;
247 }
248};
249
258template <typename IndexType = size_t, typename DistanceType = double>
260{
261 ResultItem() = default;
262 ResultItem(const IndexType index, const DistanceType distance)
263 : first(index), second(distance)
264 {
265 }
266
267 IndexType first;
268 DistanceType second;
269};
270
274template <typename _DistanceType, typename _IndexType = size_t>
276{
277 public:
278 using DistanceType = _DistanceType;
279 using IndexType = _IndexType;
280
281 public:
282 const DistanceType radius;
283
284 std::vector<ResultItem<IndexType, DistanceType>>& m_indices_dists;
285
286 explicit RadiusResultSet(
287 DistanceType radius_,
288 std::vector<ResultItem<IndexType, DistanceType>>& indices_dists)
289 : radius(radius_), m_indices_dists(indices_dists)
290 {
291 init();
292 }
293
294 void init() { clear(); }
295 void clear() { m_indices_dists.clear(); }
296
297 size_t size() const { return m_indices_dists.size(); }
298 size_t empty() const { return m_indices_dists.empty(); }
299
300 bool full() const { return true; }
301
307 bool addPoint(DistanceType dist, IndexType index)
308 {
309 if (dist < radius) m_indices_dists.emplace_back(index, dist);
310 return true;
311 }
312
313 DistanceType worstDist() const { return radius; }
314
320 {
321 if (m_indices_dists.empty())
322 throw std::runtime_error(
323 "Cannot invoke RadiusResultSet::worst_item() on "
324 "an empty list of results.");
325 auto it = std::max_element(
326 m_indices_dists.begin(), m_indices_dists.end(), IndexDist_Sorter());
327 return *it;
328 }
329};
330
335template <typename T>
336void save_value(std::ostream& stream, const T& value)
337{
338 stream.write(reinterpret_cast<const char*>(&value), sizeof(T));
339}
340
341template <typename T>
342void save_value(std::ostream& stream, const std::vector<T>& value)
343{
344 size_t size = value.size();
345 stream.write(reinterpret_cast<const char*>(&size), sizeof(size_t));
346 stream.write(reinterpret_cast<const char*>(value.data()), sizeof(T) * size);
347}
348
349template <typename T>
350void load_value(std::istream& stream, T& value)
351{
352 stream.read(reinterpret_cast<char*>(&value), sizeof(T));
353}
354
355template <typename T>
356void load_value(std::istream& stream, std::vector<T>& value)
357{
358 size_t size;
359 stream.read(reinterpret_cast<char*>(&size), sizeof(size_t));
360 value.resize(size);
361 stream.read(reinterpret_cast<char*>(value.data()), sizeof(T) * size);
362}
368struct Metric
369{
370};
371
382template <
383 class T, class DataSource, typename _DistanceType = T,
384 typename IndexType = uint32_t>
386{
387 using ElementType = T;
388 using DistanceType = _DistanceType;
389
390 const DataSource& data_source;
391
392 L1_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
393
394 DistanceType evalMetric(
395 const T* a, const IndexType b_idx, size_t size,
396 DistanceType worst_dist = -1) const
397 {
398 DistanceType result = DistanceType();
399 const T* last = a + size;
400 const T* lastgroup = last - 3;
401 size_t d = 0;
402
403 /* Process 4 items with each loop for efficiency. */
404 while (a < lastgroup)
405 {
406 const DistanceType diff0 =
407 std::abs(a[0] - data_source.kdtree_get_pt(b_idx, d++));
408 const DistanceType diff1 =
409 std::abs(a[1] - data_source.kdtree_get_pt(b_idx, d++));
410 const DistanceType diff2 =
411 std::abs(a[2] - data_source.kdtree_get_pt(b_idx, d++));
412 const DistanceType diff3 =
413 std::abs(a[3] - data_source.kdtree_get_pt(b_idx, d++));
414 result += diff0 + diff1 + diff2 + diff3;
415 a += 4;
416 if ((worst_dist > 0) && (result > worst_dist)) { return result; }
417 }
418 /* Process last 0-3 components. Not needed for standard vector lengths.
419 */
420 while (a < last)
421 {
422 result += std::abs(*a++ - data_source.kdtree_get_pt(b_idx, d++));
423 }
424 return result;
425 }
426
427 template <typename U, typename V>
428 DistanceType accum_dist(const U a, const V b, const size_t) const
429 {
430 return std::abs(a - b);
431 }
432};
433
444template <
445 class T, class DataSource, typename _DistanceType = T,
446 typename IndexType = uint32_t>
448{
449 using ElementType = T;
450 using DistanceType = _DistanceType;
451
452 const DataSource& data_source;
453
454 L2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
455
456 DistanceType evalMetric(
457 const T* a, const IndexType b_idx, size_t size,
458 DistanceType worst_dist = -1) const
459 {
460 DistanceType result = DistanceType();
461 const T* last = a + size;
462 const T* lastgroup = last - 3;
463 size_t d = 0;
464
465 /* Process 4 items with each loop for efficiency. */
466 while (a < lastgroup)
467 {
468 const DistanceType diff0 =
469 a[0] - data_source.kdtree_get_pt(b_idx, d++);
470 const DistanceType diff1 =
471 a[1] - data_source.kdtree_get_pt(b_idx, d++);
472 const DistanceType diff2 =
473 a[2] - data_source.kdtree_get_pt(b_idx, d++);
474 const DistanceType diff3 =
475 a[3] - data_source.kdtree_get_pt(b_idx, d++);
476 result +=
477 diff0 * diff0 + diff1 * diff1 + diff2 * diff2 + diff3 * diff3;
478 a += 4;
479 if ((worst_dist > 0) && (result > worst_dist)) { return result; }
480 }
481 /* Process last 0-3 components. Not needed for standard vector lengths.
482 */
483 while (a < last)
484 {
485 const DistanceType diff0 =
486 *a++ - data_source.kdtree_get_pt(b_idx, d++);
487 result += diff0 * diff0;
488 }
489 return result;
490 }
491
492 template <typename U, typename V>
493 DistanceType accum_dist(const U a, const V b, const size_t) const
494 {
495 return (a - b) * (a - b);
496 }
497};
498
509template <
510 class T, class DataSource, typename _DistanceType = T,
511 typename IndexType = uint32_t>
513{
514 using ElementType = T;
515 using DistanceType = _DistanceType;
516
517 const DataSource& data_source;
518
519 L2_Simple_Adaptor(const DataSource& _data_source)
520 : data_source(_data_source)
521 {
522 }
523
524 DistanceType evalMetric(
525 const T* a, const IndexType b_idx, size_t size) const
526 {
527 DistanceType result = DistanceType();
528 for (size_t i = 0; i < size; ++i)
529 {
530 const DistanceType diff =
531 a[i] - data_source.kdtree_get_pt(b_idx, i);
532 result += diff * diff;
533 }
534 return result;
535 }
536
537 template <typename U, typename V>
538 DistanceType accum_dist(const U a, const V b, const size_t) const
539 {
540 return (a - b) * (a - b);
541 }
542};
543
554template <
555 class T, class DataSource, typename _DistanceType = T,
556 typename IndexType = uint32_t>
558{
559 using ElementType = T;
560 using DistanceType = _DistanceType;
561
562 const DataSource& data_source;
563
564 SO2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}
565
566 DistanceType evalMetric(
567 const T* a, const IndexType b_idx, size_t size) const
568 {
569 return accum_dist(
570 a[size - 1], data_source.kdtree_get_pt(b_idx, size - 1), size - 1);
571 }
572
575 template <typename U, typename V>
576 DistanceType accum_dist(const U a, const V b, const size_t) const
577 {
578 DistanceType result = DistanceType();
579 DistanceType PI = pi_const<DistanceType>();
580 result = b - a;
581 if (result > PI)
582 result -= 2 * PI;
583 else if (result < -PI)
584 result += 2 * PI;
585 return result;
586 }
587};
588
599template <
600 class T, class DataSource, typename _DistanceType = T,
601 typename IndexType = uint32_t>
603{
604 using ElementType = T;
605 using DistanceType = _DistanceType;
606
608 distance_L2_Simple;
609
610 SO3_Adaptor(const DataSource& _data_source)
611 : distance_L2_Simple(_data_source)
612 {
613 }
614
615 DistanceType evalMetric(
616 const T* a, const IndexType b_idx, size_t size) const
617 {
618 return distance_L2_Simple.evalMetric(a, b_idx, size);
619 }
620
621 template <typename U, typename V>
622 DistanceType accum_dist(const U a, const V b, const size_t idx) const
623 {
624 return distance_L2_Simple.accum_dist(a, b, idx);
625 }
626};
627
629struct metric_L1 : public Metric
630{
631 template <class T, class DataSource, typename IndexType = uint32_t>
632 struct traits
633 {
635 };
636};
639struct metric_L2 : public Metric
640{
641 template <class T, class DataSource, typename IndexType = uint32_t>
642 struct traits
643 {
645 };
646};
650{
651 template <class T, class DataSource, typename IndexType = uint32_t>
652 struct traits
653 {
655 };
656};
658struct metric_SO2 : public Metric
659{
660 template <class T, class DataSource, typename IndexType = uint32_t>
661 struct traits
662 {
664 };
665};
667struct metric_SO3 : public Metric
668{
669 template <class T, class DataSource, typename IndexType = uint32_t>
670 struct traits
671 {
673 };
674};
675
681enum class KDTreeSingleIndexAdaptorFlags
682{
683 None = 0,
684 SkipInitialBuildIndex = 1
685};
686
687inline std::underlying_type<KDTreeSingleIndexAdaptorFlags>::type operator&(
688 KDTreeSingleIndexAdaptorFlags lhs, KDTreeSingleIndexAdaptorFlags rhs)
689{
690 using underlying =
691 typename std::underlying_type<KDTreeSingleIndexAdaptorFlags>::type;
692 return static_cast<underlying>(lhs) & static_cast<underlying>(rhs);
693}
694
697{
699 size_t _leaf_max_size = 10,
700 KDTreeSingleIndexAdaptorFlags _flags =
701 KDTreeSingleIndexAdaptorFlags::None,
702 unsigned int _n_thread_build = 1)
703 : leaf_max_size(_leaf_max_size),
704 flags(_flags),
705 n_thread_build(_n_thread_build)
706 {
707 }
708
709 size_t leaf_max_size;
710 KDTreeSingleIndexAdaptorFlags flags;
711 unsigned int n_thread_build;
712};
713
716{
717 SearchParameters(float eps_ = 0, bool sorted_ = true)
718 : eps(eps_), sorted(sorted_)
719 {
720 }
721
722 float eps;
723 bool sorted;
725};
746{
747 static constexpr size_t WORDSIZE = 16;
748 static constexpr size_t BLOCKSIZE = 8192;
749
750 /* We maintain memory alignment to word boundaries by requiring that all
751 allocations be in multiples of the machine wordsize. */
752 /* Size of machine word in bytes. Must be power of 2. */
753 /* Minimum number of bytes requested at a time from the system. Must be
754 * multiple of WORDSIZE. */
755
756 using Offset = uint32_t;
757 using Size = uint32_t;
758 using Dimension = int32_t;
759
760 Size remaining_ = 0;
761 void* base_ = nullptr;
762 void* loc_ = nullptr;
763
764 void internal_init()
765 {
766 remaining_ = 0;
767 base_ = nullptr;
768 usedMemory = 0;
769 wastedMemory = 0;
770 }
771
772 public:
773 Size usedMemory = 0;
774 Size wastedMemory = 0;
775
779 PooledAllocator() { internal_init(); }
780
784 ~PooledAllocator() { free_all(); }
785
787 void free_all()
788 {
789 while (base_ != nullptr)
790 {
791 // Get pointer to prev block
792 void* prev = *(static_cast<void**>(base_));
793 ::free(base_);
794 base_ = prev;
795 }
796 internal_init();
797 }
798
803 void* malloc(const size_t req_size)
804 {
805 /* Round size up to a multiple of wordsize. The following expression
806 only works for WORDSIZE that is a power of 2, by masking last bits
807 of incremented size to zero.
808 */
809 const Size size = (req_size + (WORDSIZE - 1)) & ~(WORDSIZE - 1);
810
811 /* Check whether a new block must be allocated. Note that the first
812 word of a block is reserved for a pointer to the previous block.
813 */
814 if (size > remaining_)
815 {
816 wastedMemory += remaining_;
817
818 /* Allocate new storage. */
819 const Size blocksize =
820 (size + sizeof(void*) + (WORDSIZE - 1) > BLOCKSIZE)
821 ? size + sizeof(void*) + (WORDSIZE - 1)
822 : BLOCKSIZE;
823
824 // use the standard C malloc to allocate memory
825 void* m = ::malloc(blocksize);
826 if (!m)
827 {
828 fprintf(stderr, "Failed to allocate memory.\n");
829 throw std::bad_alloc();
830 }
831
832 /* Fill first word of new block with pointer to previous block. */
833 static_cast<void**>(m)[0] = base_;
834 base_ = m;
835
836 Size shift = 0;
837 // int size_t = (WORDSIZE - ( (((size_t)m) + sizeof(void*)) &
838 // (WORDSIZE-1))) & (WORDSIZE-1);
839
840 remaining_ = blocksize - sizeof(void*) - shift;
841 loc_ = (static_cast<char*>(m) + sizeof(void*) + shift);
842 }
843 void* rloc = loc_;
844 loc_ = static_cast<char*>(loc_) + size;
845 remaining_ -= size;
846
847 usedMemory += size;
848
849 return rloc;
850 }
851
859 template <typename T>
860 T* allocate(const size_t count = 1)
861 {
862 T* mem = static_cast<T*>(this->malloc(sizeof(T) * count));
863 return mem;
864 }
865};
874template <int32_t DIM, typename T>
876{
877 using type = std::array<T, DIM>;
878};
880template <typename T>
881struct array_or_vector<-1, T>
882{
883 using type = std::vector<T>;
884};
885
902template <
903 class Derived, typename Distance, class DatasetAdaptor, int32_t DIM = -1,
904 typename IndexType = uint32_t>
906{
907 public:
910 void freeIndex(Derived& obj)
911 {
912 obj.pool_.free_all();
913 obj.root_node_ = nullptr;
914 obj.size_at_index_build_ = 0;
915 }
916
917 using ElementType = typename Distance::ElementType;
918 using DistanceType = typename Distance::DistanceType;
919
923 std::vector<IndexType> vAcc_;
924
925 using Offset = typename decltype(vAcc_)::size_type;
926 using Size = typename decltype(vAcc_)::size_type;
927 using Dimension = int32_t;
928
929 /*---------------------------
930 * Internal Data Structures
931 * --------------------------*/
932 struct Node
933 {
936 union
937 {
938 struct leaf
939 {
940 Offset left, right;
941 } lr;
942 struct nonleaf
943 {
944 Dimension divfeat;
946 DistanceType divlow, divhigh;
947 } sub;
948 } node_type;
949
951 Node *child1 = nullptr, *child2 = nullptr;
952 };
953
954 using NodePtr = Node*;
955 using NodeConstPtr = const Node*;
956
957 struct Interval
958 {
959 ElementType low, high;
960 };
961
962 NodePtr root_node_ = nullptr;
963
964 Size leaf_max_size_ = 0;
965
967 Size n_thread_build_ = 1;
969 Size size_ = 0;
971 Size size_at_index_build_ = 0;
972 Dimension dim_ = 0;
973
977
980 using distance_vector_t = typename array_or_vector<DIM, DistanceType>::type;
981
984
993
995 Size size(const Derived& obj) const { return obj.size_; }
996
998 Size veclen(const Derived& obj) { return DIM > 0 ? DIM : obj.dim; }
999
1001 ElementType dataset_get(
1002 const Derived& obj, IndexType element, Dimension component) const
1003 {
1004 return obj.dataset_.kdtree_get_pt(element, component);
1005 }
1006
1011 Size usedMemory(Derived& obj)
1012 {
1013 return obj.pool_.usedMemory + obj.pool_.wastedMemory +
1014 obj.dataset_.kdtree_get_point_count() *
1015 sizeof(IndexType); // pool memory and vind array memory
1016 }
1017
1018 void computeMinMax(
1019 const Derived& obj, Offset ind, Size count, Dimension element,
1020 ElementType& min_elem, ElementType& max_elem)
1021 {
1022 min_elem = dataset_get(obj, vAcc_[ind], element);
1023 max_elem = min_elem;
1024 for (Offset i = 1; i < count; ++i)
1025 {
1026 ElementType val = dataset_get(obj, vAcc_[ind + i], element);
1027 if (val < min_elem) min_elem = val;
1028 if (val > max_elem) max_elem = val;
1029 }
1030 }
1031
1040 Derived& obj, const Offset left, const Offset right, BoundingBox& bbox)
1041 {
1042 NodePtr node = obj.pool_.template allocate<Node>(); // allocate memory
1043 const auto dims = (DIM > 0 ? DIM : obj.dim_);
1044
1045 /* If too few exemplars remain, then make this a leaf node. */
1046 if ((right - left) <= static_cast<Offset>(obj.leaf_max_size_))
1047 {
1048 node->child1 = node->child2 = nullptr; /* Mark as leaf node. */
1049 node->node_type.lr.left = left;
1050 node->node_type.lr.right = right;
1051
1052 // compute bounding-box of leaf points
1053 for (Dimension i = 0; i < dims; ++i)
1054 {
1055 bbox[i].low = dataset_get(obj, obj.vAcc_[left], i);
1056 bbox[i].high = dataset_get(obj, obj.vAcc_[left], i);
1057 }
1058 for (Offset k = left + 1; k < right; ++k)
1059 {
1060 for (Dimension i = 0; i < dims; ++i)
1061 {
1062 const auto val = dataset_get(obj, obj.vAcc_[k], i);
1063 if (bbox[i].low > val) bbox[i].low = val;
1064 if (bbox[i].high < val) bbox[i].high = val;
1065 }
1066 }
1067 }
1068 else
1069 {
1070 Offset idx;
1071 Dimension cutfeat;
1072 DistanceType cutval;
1073 middleSplit_(obj, left, right - left, idx, cutfeat, cutval, bbox);
1074
1075 node->node_type.sub.divfeat = cutfeat;
1076
1077 BoundingBox left_bbox(bbox);
1078 left_bbox[cutfeat].high = cutval;
1079 node->child1 = this->divideTree(obj, left, left + idx, left_bbox);
1080
1081 BoundingBox right_bbox(bbox);
1082 right_bbox[cutfeat].low = cutval;
1083 node->child2 = this->divideTree(obj, left + idx, right, right_bbox);
1084
1085 node->node_type.sub.divlow = left_bbox[cutfeat].high;
1086 node->node_type.sub.divhigh = right_bbox[cutfeat].low;
1087
1088 for (Dimension i = 0; i < dims; ++i)
1089 {
1090 bbox[i].low = std::min(left_bbox[i].low, right_bbox[i].low);
1091 bbox[i].high = std::max(left_bbox[i].high, right_bbox[i].high);
1092 }
1093 }
1094
1095 return node;
1096 }
1097
1109 Derived& obj, const Offset left, const Offset right, BoundingBox& bbox,
1110 std::atomic<unsigned int>& thread_count, std::mutex& mutex)
1111 {
1112 std::unique_lock<std::mutex> lock(mutex);
1113 NodePtr node = obj.pool_.template allocate<Node>(); // allocate memory
1114 lock.unlock();
1115
1116 const auto dims = (DIM > 0 ? DIM : obj.dim_);
1117
1118 /* If too few exemplars remain, then make this a leaf node. */
1119 if ((right - left) <= static_cast<Offset>(obj.leaf_max_size_))
1120 {
1121 node->child1 = node->child2 = nullptr; /* Mark as leaf node. */
1122 node->node_type.lr.left = left;
1123 node->node_type.lr.right = right;
1124
1125 // compute bounding-box of leaf points
1126 for (Dimension i = 0; i < dims; ++i)
1127 {
1128 bbox[i].low = dataset_get(obj, obj.vAcc_[left], i);
1129 bbox[i].high = dataset_get(obj, obj.vAcc_[left], i);
1130 }
1131 for (Offset k = left + 1; k < right; ++k)
1132 {
1133 for (Dimension i = 0; i < dims; ++i)
1134 {
1135 const auto val = dataset_get(obj, obj.vAcc_[k], i);
1136 if (bbox[i].low > val) bbox[i].low = val;
1137 if (bbox[i].high < val) bbox[i].high = val;
1138 }
1139 }
1140 }
1141 else
1142 {
1143 Offset idx;
1144 Dimension cutfeat;
1145 DistanceType cutval;
1146 middleSplit_(obj, left, right - left, idx, cutfeat, cutval, bbox);
1147
1148 node->node_type.sub.divfeat = cutfeat;
1149
1150 std::future<NodePtr> left_future, right_future;
1151
1152 BoundingBox left_bbox(bbox);
1153 left_bbox[cutfeat].high = cutval;
1154 if (++thread_count < n_thread_build_)
1155 {
1156 left_future = std::async(
1157 std::launch::async, &KDTreeBaseClass::divideTreeConcurrent,
1158 this, std::ref(obj), left, left + idx, std::ref(left_bbox),
1159 std::ref(thread_count), std::ref(mutex));
1160 }
1161 else
1162 {
1163 --thread_count;
1164 node->child1 = this->divideTreeConcurrent(
1165 obj, left, left + idx, left_bbox, thread_count, mutex);
1166 }
1167
1168 BoundingBox right_bbox(bbox);
1169 right_bbox[cutfeat].low = cutval;
1170 if (++thread_count < n_thread_build_)
1171 {
1172 right_future = std::async(
1173 std::launch::async, &KDTreeBaseClass::divideTreeConcurrent,
1174 this, std::ref(obj), left + idx, right,
1175 std::ref(right_bbox), std::ref(thread_count),
1176 std::ref(mutex));
1177 }
1178 else
1179 {
1180 --thread_count;
1181 node->child2 = this->divideTreeConcurrent(
1182 obj, left + idx, right, right_bbox, thread_count, mutex);
1183 }
1184
1185 if (left_future.valid())
1186 {
1187 node->child1 = left_future.get();
1188 --thread_count;
1189 }
1190 if (right_future.valid())
1191 {
1192 node->child2 = right_future.get();
1193 --thread_count;
1194 }
1195
1196 node->node_type.sub.divlow = left_bbox[cutfeat].high;
1197 node->node_type.sub.divhigh = right_bbox[cutfeat].low;
1198
1199 for (Dimension i = 0; i < dims; ++i)
1200 {
1201 bbox[i].low = std::min(left_bbox[i].low, right_bbox[i].low);
1202 bbox[i].high = std::max(left_bbox[i].high, right_bbox[i].high);
1203 }
1204 }
1205
1206 return node;
1207 }
1208
1209 void middleSplit_(
1210 const Derived& obj, const Offset ind, const Size count, Offset& index,
1211 Dimension& cutfeat, DistanceType& cutval, const BoundingBox& bbox)
1212 {
1213 const auto dims = (DIM > 0 ? DIM : obj.dim_);
1214 const auto EPS = static_cast<DistanceType>(0.00001);
1215 ElementType max_span = bbox[0].high - bbox[0].low;
1216 for (Dimension i = 1; i < dims; ++i)
1217 {
1218 ElementType span = bbox[i].high - bbox[i].low;
1219 if (span > max_span) { max_span = span; }
1220 }
1221 ElementType max_spread = -1;
1222 cutfeat = 0;
1223 for (Dimension i = 0; i < dims; ++i)
1224 {
1225 ElementType span = bbox[i].high - bbox[i].low;
1226 if (span > (1 - EPS) * max_span)
1227 {
1228 ElementType min_elem, max_elem;
1229 computeMinMax(obj, ind, count, i, min_elem, max_elem);
1230 ElementType spread = max_elem - min_elem;
1231 if (spread > max_spread)
1232 {
1233 cutfeat = i;
1234 max_spread = spread;
1235 }
1236 }
1237 }
1238 // split in the middle
1239 DistanceType split_val = (bbox[cutfeat].low + bbox[cutfeat].high) / 2;
1240 ElementType min_elem, max_elem;
1241 computeMinMax(obj, ind, count, cutfeat, min_elem, max_elem);
1242
1243 if (split_val < min_elem)
1244 cutval = min_elem;
1245 else if (split_val > max_elem)
1246 cutval = max_elem;
1247 else
1248 cutval = split_val;
1249
1250 Offset lim1, lim2;
1251 planeSplit(obj, ind, count, cutfeat, cutval, lim1, lim2);
1252
1253 if (lim1 > count / 2)
1254 index = lim1;
1255 else if (lim2 < count / 2)
1256 index = lim2;
1257 else
1258 index = count / 2;
1259 }
1260
1271 const Derived& obj, const Offset ind, const Size count,
1272 const Dimension cutfeat, const DistanceType& cutval, Offset& lim1,
1273 Offset& lim2)
1274 {
1275 /* Move vector indices for left subtree to front of list. */
1276 Offset left = 0;
1277 Offset right = count - 1;
1278 for (;;)
1279 {
1280 while (left <= right &&
1281 dataset_get(obj, vAcc_[ind + left], cutfeat) < cutval)
1282 ++left;
1283 while (right && left <= right &&
1284 dataset_get(obj, vAcc_[ind + right], cutfeat) >= cutval)
1285 --right;
1286 if (left > right || !right)
1287 break; // "!right" was added to support unsigned Index types
1288 std::swap(vAcc_[ind + left], vAcc_[ind + right]);
1289 ++left;
1290 --right;
1291 }
1292 /* If either list is empty, it means that all remaining features
1293 * are identical. Split in the middle to maintain a balanced tree.
1294 */
1295 lim1 = left;
1296 right = count - 1;
1297 for (;;)
1298 {
1299 while (left <= right &&
1300 dataset_get(obj, vAcc_[ind + left], cutfeat) <= cutval)
1301 ++left;
1302 while (right && left <= right &&
1303 dataset_get(obj, vAcc_[ind + right], cutfeat) > cutval)
1304 --right;
1305 if (left > right || !right)
1306 break; // "!right" was added to support unsigned Index types
1307 std::swap(vAcc_[ind + left], vAcc_[ind + right]);
1308 ++left;
1309 --right;
1310 }
1311 lim2 = left;
1312 }
1313
1314 DistanceType computeInitialDistances(
1315 const Derived& obj, const ElementType* vec,
1316 distance_vector_t& dists) const
1317 {
1318 assert(vec);
1319 DistanceType dist = DistanceType();
1320
1321 for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim_); ++i)
1322 {
1323 if (vec[i] < obj.root_bbox_[i].low)
1324 {
1325 dists[i] =
1326 obj.distance_.accum_dist(vec[i], obj.root_bbox_[i].low, i);
1327 dist += dists[i];
1328 }
1329 if (vec[i] > obj.root_bbox_[i].high)
1330 {
1331 dists[i] =
1332 obj.distance_.accum_dist(vec[i], obj.root_bbox_[i].high, i);
1333 dist += dists[i];
1334 }
1335 }
1336 return dist;
1337 }
1338
1339 static void save_tree(
1340 const Derived& obj, std::ostream& stream, const NodeConstPtr tree)
1341 {
1342 save_value(stream, *tree);
1343 if (tree->child1 != nullptr) { save_tree(obj, stream, tree->child1); }
1344 if (tree->child2 != nullptr) { save_tree(obj, stream, tree->child2); }
1345 }
1346
1347 static void load_tree(Derived& obj, std::istream& stream, NodePtr& tree)
1348 {
1349 tree = obj.pool_.template allocate<Node>();
1350 load_value(stream, *tree);
1351 if (tree->child1 != nullptr) { load_tree(obj, stream, tree->child1); }
1352 if (tree->child2 != nullptr) { load_tree(obj, stream, tree->child2); }
1353 }
1354
1360 void saveIndex(const Derived& obj, std::ostream& stream) const
1361 {
1362 save_value(stream, obj.size_);
1363 save_value(stream, obj.dim_);
1364 save_value(stream, obj.root_bbox_);
1365 save_value(stream, obj.leaf_max_size_);
1366 save_value(stream, obj.vAcc_);
1367 if (obj.root_node_) save_tree(obj, stream, obj.root_node_);
1368 }
1369
1375 void loadIndex(Derived& obj, std::istream& stream)
1376 {
1377 load_value(stream, obj.size_);
1378 load_value(stream, obj.dim_);
1379 load_value(stream, obj.root_bbox_);
1380 load_value(stream, obj.leaf_max_size_);
1381 load_value(stream, obj.vAcc_);
1382 load_tree(obj, stream, obj.root_node_);
1383 }
1384};
1385
1427template <
1428 typename Distance, class DatasetAdaptor, int32_t DIM = -1,
1429 typename IndexType = uint32_t>
1431 : public KDTreeBaseClass<
1432 KDTreeSingleIndexAdaptor<Distance, DatasetAdaptor, DIM, IndexType>,
1433 Distance, DatasetAdaptor, DIM, IndexType>
1434{
1435 public:
1439 Distance, DatasetAdaptor, DIM, IndexType>&) = delete;
1440
1442 const DatasetAdaptor& dataset_;
1443
1444 const KDTreeSingleIndexAdaptorParams indexParams;
1445
1446 Distance distance_;
1447
1448 using Base = typename nanoflann::KDTreeBaseClass<
1450 Distance, DatasetAdaptor, DIM, IndexType>,
1451 Distance, DatasetAdaptor, DIM, IndexType>;
1452
1453 using Offset = typename Base::Offset;
1454 using Size = typename Base::Size;
1455 using Dimension = typename Base::Dimension;
1456
1457 using ElementType = typename Base::ElementType;
1458 using DistanceType = typename Base::DistanceType;
1459
1460 using Node = typename Base::Node;
1461 using NodePtr = Node*;
1462
1463 using Interval = typename Base::Interval;
1464
1467 using BoundingBox = typename Base::BoundingBox;
1468
1471 using distance_vector_t = typename Base::distance_vector_t;
1472
1493 template <class... Args>
1495 const Dimension dimensionality, const DatasetAdaptor& inputData,
1496 const KDTreeSingleIndexAdaptorParams& params, Args&&... args)
1497 : dataset_(inputData),
1498 indexParams(params),
1499 distance_(inputData, std::forward<Args>(args)...)
1500 {
1501 init(dimensionality, params);
1502 }
1503
1504 explicit KDTreeSingleIndexAdaptor(
1505 const Dimension dimensionality, const DatasetAdaptor& inputData,
1506 const KDTreeSingleIndexAdaptorParams& params = {})
1507 : dataset_(inputData), indexParams(params), distance_(inputData)
1508 {
1509 init(dimensionality, params);
1510 }
1511
1512 private:
1513 void init(
1514 const Dimension dimensionality,
1515 const KDTreeSingleIndexAdaptorParams& params)
1516 {
1517 Base::size_ = dataset_.kdtree_get_point_count();
1518 Base::size_at_index_build_ = Base::size_;
1519 Base::dim_ = dimensionality;
1520 if (DIM > 0) Base::dim_ = DIM;
1521 Base::leaf_max_size_ = params.leaf_max_size;
1522 if (params.n_thread_build > 0)
1523 {
1524 Base::n_thread_build_ = params.n_thread_build;
1525 }
1526 else
1527 {
1528 Base::n_thread_build_ =
1529 std::max(std::thread::hardware_concurrency(), 1u);
1530 }
1531
1532 if (!(params.flags &
1533 KDTreeSingleIndexAdaptorFlags::SkipInitialBuildIndex))
1534 {
1535 // Build KD-tree:
1536 buildIndex();
1537 }
1538 }
1539
1540 public:
1545 {
1546 Base::size_ = dataset_.kdtree_get_point_count();
1547 Base::size_at_index_build_ = Base::size_;
1548 init_vind();
1549 this->freeIndex(*this);
1550 Base::size_at_index_build_ = Base::size_;
1551 if (Base::size_ == 0) return;
1552 computeBoundingBox(Base::root_bbox_);
1553 // construct the tree
1554 if (Base::n_thread_build_ == 1)
1555 {
1556 Base::root_node_ =
1557 this->divideTree(*this, 0, Base::size_, Base::root_bbox_);
1558 }
1559 else
1560 {
1561 std::atomic<unsigned int> thread_count(0u);
1562 std::mutex mutex;
1563 Base::root_node_ = this->divideTreeConcurrent(
1564 *this, 0, Base::size_, Base::root_bbox_, thread_count, mutex);
1565 }
1566 }
1567
1587 template <typename RESULTSET>
1589 RESULTSET& result, const ElementType* vec,
1590 const SearchParameters& searchParams = {}) const
1591 {
1592 assert(vec);
1593 if (this->size(*this) == 0) return false;
1594 if (!Base::root_node_)
1595 throw std::runtime_error(
1596 "[nanoflann] findNeighbors() called before building the "
1597 "index.");
1598 float epsError = 1 + searchParams.eps;
1599
1600 // fixed or variable-sized container (depending on DIM)
1601 distance_vector_t dists;
1602 // Fill it with zeros.
1603 auto zero = static_cast<decltype(result.worstDist())>(0);
1604 assign(dists, (DIM > 0 ? DIM : Base::dim_), zero);
1605 DistanceType dist = this->computeInitialDistances(*this, vec, dists);
1606 searchLevel(result, vec, Base::root_node_, dist, dists, epsError);
1607 return result.full();
1608 }
1609
1626 const ElementType* query_point, const Size num_closest,
1627 IndexType* out_indices, DistanceType* out_distances) const
1628 {
1630 resultSet.init(out_indices, out_distances);
1631 findNeighbors(resultSet, query_point);
1632 return resultSet.size();
1633 }
1634
1655 const ElementType* query_point, const DistanceType& radius,
1656 std::vector<ResultItem<IndexType, DistanceType>>& IndicesDists,
1657 const SearchParameters& searchParams = {}) const
1658 {
1660 radius, IndicesDists);
1661 const Size nFound =
1662 radiusSearchCustomCallback(query_point, resultSet, searchParams);
1663 if (searchParams.sorted)
1664 std::sort(
1665 IndicesDists.begin(), IndicesDists.end(), IndexDist_Sorter());
1666 return nFound;
1667 }
1668
1674 template <class SEARCH_CALLBACK>
1676 const ElementType* query_point, SEARCH_CALLBACK& resultSet,
1677 const SearchParameters& searchParams = {}) const
1678 {
1679 findNeighbors(resultSet, query_point, searchParams);
1680 return resultSet.size();
1681 }
1682
1685 public:
1689 {
1690 // Create a permutable array of indices to the input vectors.
1691 Base::size_ = dataset_.kdtree_get_point_count();
1692 if (Base::vAcc_.size() != Base::size_) Base::vAcc_.resize(Base::size_);
1693 for (Size i = 0; i < Base::size_; i++) Base::vAcc_[i] = i;
1694 }
1695
1696 void computeBoundingBox(BoundingBox& bbox)
1697 {
1698 const auto dims = (DIM > 0 ? DIM : Base::dim_);
1699 resize(bbox, dims);
1700 if (dataset_.kdtree_get_bbox(bbox))
1701 {
1702 // Done! It was implemented in derived class
1703 }
1704 else
1705 {
1706 const Size N = dataset_.kdtree_get_point_count();
1707 if (!N)
1708 throw std::runtime_error(
1709 "[nanoflann] computeBoundingBox() called but "
1710 "no data points found.");
1711 for (Dimension i = 0; i < dims; ++i)
1712 {
1713 bbox[i].low = bbox[i].high =
1714 this->dataset_get(*this, Base::vAcc_[0], i);
1715 }
1716 for (Offset k = 1; k < N; ++k)
1717 {
1718 for (Dimension i = 0; i < dims; ++i)
1719 {
1720 const auto val =
1721 this->dataset_get(*this, Base::vAcc_[k], i);
1722 if (val < bbox[i].low) bbox[i].low = val;
1723 if (val > bbox[i].high) bbox[i].high = val;
1724 }
1725 }
1726 }
1727 }
1728
1735 template <class RESULTSET>
1737 RESULTSET& result_set, const ElementType* vec, const NodePtr node,
1738 DistanceType mindist, distance_vector_t& dists,
1739 const float epsError) const
1740 {
1741 /* If this is a leaf node, then do check and return. */
1742 if ((node->child1 == nullptr) && (node->child2 == nullptr))
1743 {
1744 DistanceType worst_dist = result_set.worstDist();
1745 for (Offset i = node->node_type.lr.left;
1746 i < node->node_type.lr.right; ++i)
1747 {
1748 const IndexType accessor = Base::vAcc_[i]; // reorder... : i;
1749 DistanceType dist = distance_.evalMetric(
1750 vec, accessor, (DIM > 0 ? DIM : Base::dim_));
1751 if (dist < worst_dist)
1752 {
1753 if (!result_set.addPoint(dist, Base::vAcc_[i]))
1754 {
1755 // the resultset doesn't want to receive any more
1756 // points, we're done searching!
1757 return false;
1758 }
1759 }
1760 }
1761 return true;
1762 }
1763
1764 /* Which child branch should be taken first? */
1765 Dimension idx = node->node_type.sub.divfeat;
1766 ElementType val = vec[idx];
1767 DistanceType diff1 = val - node->node_type.sub.divlow;
1768 DistanceType diff2 = val - node->node_type.sub.divhigh;
1769
1770 NodePtr bestChild;
1771 NodePtr otherChild;
1772 DistanceType cut_dist;
1773 if ((diff1 + diff2) < 0)
1774 {
1775 bestChild = node->child1;
1776 otherChild = node->child2;
1777 cut_dist =
1778 distance_.accum_dist(val, node->node_type.sub.divhigh, idx);
1779 }
1780 else
1781 {
1782 bestChild = node->child2;
1783 otherChild = node->child1;
1784 cut_dist =
1785 distance_.accum_dist(val, node->node_type.sub.divlow, idx);
1786 }
1787
1788 /* Call recursively to search next level down. */
1789 if (!searchLevel(result_set, vec, bestChild, mindist, dists, epsError))
1790 {
1791 // the resultset doesn't want to receive any more points, we're done
1792 // searching!
1793 return false;
1794 }
1795
1796 DistanceType dst = dists[idx];
1797 mindist = mindist + cut_dist - dst;
1798 dists[idx] = cut_dist;
1799 if (mindist * epsError <= result_set.worstDist())
1800 {
1801 if (!searchLevel(
1802 result_set, vec, otherChild, mindist, dists, epsError))
1803 {
1804 // the resultset doesn't want to receive any more points, we're
1805 // done searching!
1806 return false;
1807 }
1808 }
1809 dists[idx] = dst;
1810 return true;
1811 }
1812
1813 public:
1819 void saveIndex(std::ostream& stream) const
1820 {
1821 Base::saveIndex(*this, stream);
1822 }
1823
1829 void loadIndex(std::istream& stream) { Base::loadIndex(*this, stream); }
1830
1831}; // class KDTree
1832
1870template <
1871 typename Distance, class DatasetAdaptor, int32_t DIM = -1,
1872 typename IndexType = uint32_t>
1874 : public KDTreeBaseClass<
1875 KDTreeSingleIndexDynamicAdaptor_<
1876 Distance, DatasetAdaptor, DIM, IndexType>,
1877 Distance, DatasetAdaptor, DIM, IndexType>
1878{
1879 public:
1883 const DatasetAdaptor& dataset_;
1884
1885 KDTreeSingleIndexAdaptorParams index_params_;
1886
1887 std::vector<int>& treeIndex_;
1888
1889 Distance distance_;
1890
1891 using Base = typename nanoflann::KDTreeBaseClass<
1893 Distance, DatasetAdaptor, DIM, IndexType>,
1894 Distance, DatasetAdaptor, DIM, IndexType>;
1895
1896 using ElementType = typename Base::ElementType;
1897 using DistanceType = typename Base::DistanceType;
1898
1899 using Offset = typename Base::Offset;
1900 using Size = typename Base::Size;
1901 using Dimension = typename Base::Dimension;
1902
1903 using Node = typename Base::Node;
1904 using NodePtr = Node*;
1905
1906 using Interval = typename Base::Interval;
1909 using BoundingBox = typename Base::BoundingBox;
1910
1913 using distance_vector_t = typename Base::distance_vector_t;
1914
1931 const Dimension dimensionality, const DatasetAdaptor& inputData,
1932 std::vector<int>& treeIndex,
1933 const KDTreeSingleIndexAdaptorParams& params =
1935 : dataset_(inputData),
1936 index_params_(params),
1937 treeIndex_(treeIndex),
1938 distance_(inputData)
1939 {
1940 Base::size_ = 0;
1941 Base::size_at_index_build_ = 0;
1942 for (auto& v : Base::root_bbox_) v = {};
1943 Base::dim_ = dimensionality;
1944 if (DIM > 0) Base::dim_ = DIM;
1945 Base::leaf_max_size_ = params.leaf_max_size;
1946 if (params.n_thread_build > 0)
1947 {
1948 Base::n_thread_build_ = params.n_thread_build;
1949 }
1950 else
1951 {
1952 Base::n_thread_build_ =
1953 std::max(std::thread::hardware_concurrency(), 1u);
1954 }
1955 }
1956
1959 const KDTreeSingleIndexDynamicAdaptor_& rhs) = default;
1960
1964 {
1966 std::swap(Base::vAcc_, tmp.Base::vAcc_);
1967 std::swap(Base::leaf_max_size_, tmp.Base::leaf_max_size_);
1968 std::swap(index_params_, tmp.index_params_);
1969 std::swap(treeIndex_, tmp.treeIndex_);
1970 std::swap(Base::size_, tmp.Base::size_);
1971 std::swap(Base::size_at_index_build_, tmp.Base::size_at_index_build_);
1972 std::swap(Base::root_node_, tmp.Base::root_node_);
1973 std::swap(Base::root_bbox_, tmp.Base::root_bbox_);
1974 std::swap(Base::pool_, tmp.Base::pool_);
1975 return *this;
1976 }
1977
1982 {
1983 Base::size_ = Base::vAcc_.size();
1984 this->freeIndex(*this);
1985 Base::size_at_index_build_ = Base::size_;
1986 if (Base::size_ == 0) return;
1987 computeBoundingBox(Base::root_bbox_);
1988 // construct the tree
1989 if (Base::n_thread_build_ == 1)
1990 {
1991 Base::root_node_ =
1992 this->divideTree(*this, 0, Base::size_, Base::root_bbox_);
1993 }
1994 else
1995 {
1996 std::atomic<unsigned int> thread_count(0u);
1997 std::mutex mutex;
1998 Base::root_node_ = this->divideTreeConcurrent(
1999 *this, 0, Base::size_, Base::root_bbox_, thread_count, mutex);
2000 }
2001 }
2002
2026 template <typename RESULTSET>
2028 RESULTSET& result, const ElementType* vec,
2029 const SearchParameters& searchParams = {}) const
2030 {
2031 assert(vec);
2032 if (this->size(*this) == 0) return false;
2033 if (!Base::root_node_) return false;
2034 float epsError = 1 + searchParams.eps;
2035
2036 // fixed or variable-sized container (depending on DIM)
2037 distance_vector_t dists;
2038 // Fill it with zeros.
2039 assign(
2040 dists, (DIM > 0 ? DIM : Base::dim_),
2041 static_cast<typename distance_vector_t::value_type>(0));
2042 DistanceType dist = this->computeInitialDistances(*this, vec, dists);
2043 searchLevel(result, vec, Base::root_node_, dist, dists, epsError);
2044 return result.full();
2045 }
2046
2062 const ElementType* query_point, const Size num_closest,
2063 IndexType* out_indices, DistanceType* out_distances) const
2064 {
2066 resultSet.init(out_indices, out_distances);
2067 findNeighbors(resultSet, query_point);
2068 return resultSet.size();
2069 }
2070
2091 const ElementType* query_point, const DistanceType& radius,
2092 std::vector<ResultItem<IndexType, DistanceType>>& IndicesDists,
2093 const SearchParameters& searchParams = {}) const
2094 {
2096 radius, IndicesDists);
2097 const size_t nFound =
2098 radiusSearchCustomCallback(query_point, resultSet, searchParams);
2099 if (searchParams.sorted)
2100 std::sort(
2101 IndicesDists.begin(), IndicesDists.end(), IndexDist_Sorter());
2102 return nFound;
2103 }
2104
2110 template <class SEARCH_CALLBACK>
2112 const ElementType* query_point, SEARCH_CALLBACK& resultSet,
2113 const SearchParameters& searchParams = {}) const
2114 {
2115 findNeighbors(resultSet, query_point, searchParams);
2116 return resultSet.size();
2117 }
2118
2121 public:
2122 void computeBoundingBox(BoundingBox& bbox)
2123 {
2124 const auto dims = (DIM > 0 ? DIM : Base::dim_);
2125 resize(bbox, dims);
2126
2127 if (dataset_.kdtree_get_bbox(bbox))
2128 {
2129 // Done! It was implemented in derived class
2130 }
2131 else
2132 {
2133 const Size N = Base::size_;
2134 if (!N)
2135 throw std::runtime_error(
2136 "[nanoflann] computeBoundingBox() called but "
2137 "no data points found.");
2138 for (Dimension i = 0; i < dims; ++i)
2139 {
2140 bbox[i].low = bbox[i].high =
2141 this->dataset_get(*this, Base::vAcc_[0], i);
2142 }
2143 for (Offset k = 1; k < N; ++k)
2144 {
2145 for (Dimension i = 0; i < dims; ++i)
2146 {
2147 const auto val =
2148 this->dataset_get(*this, Base::vAcc_[k], i);
2149 if (val < bbox[i].low) bbox[i].low = val;
2150 if (val > bbox[i].high) bbox[i].high = val;
2151 }
2152 }
2153 }
2154 }
2155
2160 template <class RESULTSET>
2162 RESULTSET& result_set, const ElementType* vec, const NodePtr node,
2163 DistanceType mindist, distance_vector_t& dists,
2164 const float epsError) const
2165 {
2166 /* If this is a leaf node, then do check and return. */
2167 if ((node->child1 == nullptr) && (node->child2 == nullptr))
2168 {
2169 DistanceType worst_dist = result_set.worstDist();
2170 for (Offset i = node->node_type.lr.left;
2171 i < node->node_type.lr.right; ++i)
2172 {
2173 const IndexType index = Base::vAcc_[i]; // reorder... : i;
2174 if (treeIndex_[index] == -1) continue;
2175 DistanceType dist = distance_.evalMetric(
2176 vec, index, (DIM > 0 ? DIM : Base::dim_));
2177 if (dist < worst_dist)
2178 {
2179 if (!result_set.addPoint(
2180 static_cast<typename RESULTSET::DistanceType>(dist),
2181 static_cast<typename RESULTSET::IndexType>(
2182 Base::vAcc_[i])))
2183 {
2184 // the resultset doesn't want to receive any more
2185 // points, we're done searching!
2186 return; // false;
2187 }
2188 }
2189 }
2190 return;
2191 }
2192
2193 /* Which child branch should be taken first? */
2194 Dimension idx = node->node_type.sub.divfeat;
2195 ElementType val = vec[idx];
2196 DistanceType diff1 = val - node->node_type.sub.divlow;
2197 DistanceType diff2 = val - node->node_type.sub.divhigh;
2198
2199 NodePtr bestChild;
2200 NodePtr otherChild;
2201 DistanceType cut_dist;
2202 if ((diff1 + diff2) < 0)
2203 {
2204 bestChild = node->child1;
2205 otherChild = node->child2;
2206 cut_dist =
2207 distance_.accum_dist(val, node->node_type.sub.divhigh, idx);
2208 }
2209 else
2210 {
2211 bestChild = node->child2;
2212 otherChild = node->child1;
2213 cut_dist =
2214 distance_.accum_dist(val, node->node_type.sub.divlow, idx);
2215 }
2216
2217 /* Call recursively to search next level down. */
2218 searchLevel(result_set, vec, bestChild, mindist, dists, epsError);
2219
2220 DistanceType dst = dists[idx];
2221 mindist = mindist + cut_dist - dst;
2222 dists[idx] = cut_dist;
2223 if (mindist * epsError <= result_set.worstDist())
2224 {
2225 searchLevel(result_set, vec, otherChild, mindist, dists, epsError);
2226 }
2227 dists[idx] = dst;
2228 }
2229
2230 public:
2236 void saveIndex(std::ostream& stream) { saveIndex(*this, stream); }
2237
2243 void loadIndex(std::istream& stream) { loadIndex(*this, stream); }
2244};
2245
2260template <
2261 typename Distance, class DatasetAdaptor, int32_t DIM = -1,
2262 typename IndexType = uint32_t>
2264{
2265 public:
2266 using ElementType = typename Distance::ElementType;
2267 using DistanceType = typename Distance::DistanceType;
2268
2269 using Offset = typename KDTreeSingleIndexDynamicAdaptor_<
2270 Distance, DatasetAdaptor, DIM>::Offset;
2271 using Size = typename KDTreeSingleIndexDynamicAdaptor_<
2272 Distance, DatasetAdaptor, DIM>::Size;
2273 using Dimension = typename KDTreeSingleIndexDynamicAdaptor_<
2274 Distance, DatasetAdaptor, DIM>::Dimension;
2275
2276 protected:
2277 Size leaf_max_size_;
2278 Size treeCount_;
2279 Size pointCount_;
2280
2284 const DatasetAdaptor& dataset_;
2285
2288 std::vector<int> treeIndex_;
2289 std::unordered_set<int> removedPoints_;
2290
2291 KDTreeSingleIndexAdaptorParams index_params_;
2292
2293 Dimension dim_;
2294
2296 Distance, DatasetAdaptor, DIM, IndexType>;
2297 std::vector<index_container_t> index_;
2298
2299 public:
2302 const std::vector<index_container_t>& getAllIndices() const
2303 {
2304 return index_;
2305 }
2306
2307 private:
2309 int First0Bit(IndexType num)
2310 {
2311 int pos = 0;
2312 while (num & 1)
2313 {
2314 num = num >> 1;
2315 pos++;
2316 }
2317 return pos;
2318 }
2319
2321 void init()
2322 {
2323 using my_kd_tree_t = KDTreeSingleIndexDynamicAdaptor_<
2324 Distance, DatasetAdaptor, DIM, IndexType>;
2325 std::vector<my_kd_tree_t> index(
2326 treeCount_,
2327 my_kd_tree_t(dim_ /*dim*/, dataset_, treeIndex_, index_params_));
2328 index_ = index;
2329 }
2330
2331 public:
2332 Distance distance_;
2333
2350 const int dimensionality, const DatasetAdaptor& inputData,
2351 const KDTreeSingleIndexAdaptorParams& params =
2353 const size_t maximumPointCount = 1000000000U)
2354 : dataset_(inputData), index_params_(params), distance_(inputData)
2355 {
2356 treeCount_ = static_cast<size_t>(std::log2(maximumPointCount)) + 1;
2357 pointCount_ = 0U;
2358 dim_ = dimensionality;
2359 treeIndex_.clear();
2360 if (DIM > 0) dim_ = DIM;
2361 leaf_max_size_ = params.leaf_max_size;
2362 init();
2363 const size_t num_initial_points = dataset_.kdtree_get_point_count();
2364 if (num_initial_points > 0) { addPoints(0, num_initial_points - 1); }
2365 }
2366
2370 Distance, DatasetAdaptor, DIM, IndexType>&) = delete;
2371
2373 void addPoints(IndexType start, IndexType end)
2374 {
2375 const Size count = end - start + 1;
2376 int maxIndex = 0;
2377 treeIndex_.resize(treeIndex_.size() + count);
2378 for (IndexType idx = start; idx <= end; idx++)
2379 {
2380 const int pos = First0Bit(pointCount_);
2381 maxIndex = std::max(pos, maxIndex);
2382 treeIndex_[pointCount_] = pos;
2383
2384 const auto it = removedPoints_.find(idx);
2385 if (it != removedPoints_.end())
2386 {
2387 removedPoints_.erase(it);
2388 treeIndex_[idx] = pos;
2389 }
2390
2391 for (int i = 0; i < pos; i++)
2392 {
2393 for (int j = 0; j < static_cast<int>(index_[i].vAcc_.size());
2394 j++)
2395 {
2396 index_[pos].vAcc_.push_back(index_[i].vAcc_[j]);
2397 if (treeIndex_[index_[i].vAcc_[j]] != -1)
2398 treeIndex_[index_[i].vAcc_[j]] = pos;
2399 }
2400 index_[i].vAcc_.clear();
2401 }
2402 index_[pos].vAcc_.push_back(idx);
2403 pointCount_++;
2404 }
2405
2406 for (int i = 0; i <= maxIndex; ++i)
2407 {
2408 index_[i].freeIndex(index_[i]);
2409 if (!index_[i].vAcc_.empty()) index_[i].buildIndex();
2410 }
2411 }
2412
2414 void removePoint(size_t idx)
2415 {
2416 if (idx >= pointCount_) return;
2417 removedPoints_.insert(idx);
2418 treeIndex_[idx] = -1;
2419 }
2420
2437 template <typename RESULTSET>
2439 RESULTSET& result, const ElementType* vec,
2440 const SearchParameters& searchParams = {}) const
2441 {
2442 for (size_t i = 0; i < treeCount_; i++)
2443 {
2444 index_[i].findNeighbors(result, &vec[0], searchParams);
2445 }
2446 return result.full();
2447 }
2448};
2449
2475template <
2476 class MatrixType, int32_t DIM = -1, class Distance = nanoflann::metric_L2,
2477 bool row_major = true>
2479{
2480 using self_t =
2482 using num_t = typename MatrixType::Scalar;
2483 using IndexType = typename MatrixType::Index;
2484 using metric_t = typename Distance::template traits<
2485 num_t, self_t, IndexType>::distance_t;
2486
2488 metric_t, self_t,
2489 row_major ? MatrixType::ColsAtCompileTime
2490 : MatrixType::RowsAtCompileTime,
2491 IndexType>;
2492
2493 index_t* index_;
2495
2496 using Offset = typename index_t::Offset;
2497 using Size = typename index_t::Size;
2498 using Dimension = typename index_t::Dimension;
2499
2502 const Dimension dimensionality,
2503 const std::reference_wrapper<const MatrixType>& mat,
2504 const int leaf_max_size = 10)
2505 : m_data_matrix(mat)
2506 {
2507 const auto dims = row_major ? mat.get().cols() : mat.get().rows();
2508 if (static_cast<Dimension>(dims) != dimensionality)
2509 throw std::runtime_error(
2510 "Error: 'dimensionality' must match column count in data "
2511 "matrix");
2512 if (DIM > 0 && static_cast<int32_t>(dims) != DIM)
2513 throw std::runtime_error(
2514 "Data set dimensionality does not match the 'DIM' template "
2515 "argument");
2516 index_ = new index_t(
2517 dims, *this /* adaptor */,
2519 }
2520
2521 public:
2524
2525 ~KDTreeEigenMatrixAdaptor() { delete index_; }
2526
2527 const std::reference_wrapper<const MatrixType> m_data_matrix;
2528
2537 void query(
2538 const num_t* query_point, const Size num_closest,
2539 IndexType* out_indices, num_t* out_distances) const
2540 {
2541 nanoflann::KNNResultSet<num_t, IndexType> resultSet(num_closest);
2542 resultSet.init(out_indices, out_distances);
2543 index_->findNeighbors(resultSet, query_point);
2544 }
2545
2549 const self_t& derived() const { return *this; }
2550 self_t& derived() { return *this; }
2551
2552 // Must return the number of data points
2553 Size kdtree_get_point_count() const
2554 {
2555 if (row_major)
2556 return m_data_matrix.get().rows();
2557 else
2558 return m_data_matrix.get().cols();
2559 }
2560
2561 // Returns the dim'th component of the idx'th point in the class:
2562 num_t kdtree_get_pt(const IndexType idx, size_t dim) const
2563 {
2564 if (row_major)
2565 return m_data_matrix.get().coeff(idx, IndexType(dim));
2566 else
2567 return m_data_matrix.get().coeff(IndexType(dim), idx);
2568 }
2569
2570 // Optional bounding-box computation: return false to default to a standard
2571 // bbox computation loop.
2572 // Return true if the BBOX was already computed by the class and returned
2573 // in "bb" so it can be avoided to redo it again. Look at bb.size() to
2574 // find out the expected dimensionality (e.g. 2 or 3 for point clouds)
2575 template <class BBOX>
2576 bool kdtree_get_bbox(BBOX& /*bb*/) const
2577 {
2578 return false;
2579 }
2580
2583}; // end of KDTreeEigenMatrixAdaptor // end of grouping
2587} // namespace nanoflann
Definition: nanoflann.hpp:906
Size usedMemory(Derived &obj)
Definition: nanoflann.hpp:1011
NodePtr divideTreeConcurrent(Derived &obj, const Offset left, const Offset right, BoundingBox &bbox, std::atomic< unsigned int > &thread_count, std::mutex &mutex)
Definition: nanoflann.hpp:1108
void planeSplit(const Derived &obj, const Offset ind, const Size count, const Dimension cutfeat, const DistanceType &cutval, Offset &lim1, Offset &lim2)
Definition: nanoflann.hpp:1270
typename array_or_vector< DIM, DistanceType >::type distance_vector_t
Definition: nanoflann.hpp:980
std::vector< IndexType > vAcc_
Definition: nanoflann.hpp:923
Size size(const Derived &obj) const
Definition: nanoflann.hpp:995
void freeIndex(Derived &obj)
Definition: nanoflann.hpp:910
void loadIndex(Derived &obj, std::istream &stream)
Definition: nanoflann.hpp:1375
typename array_or_vector< DIM, Interval >::type BoundingBox
Definition: nanoflann.hpp:976
BoundingBox root_bbox_
Definition: nanoflann.hpp:983
void saveIndex(const Derived &obj, std::ostream &stream) const
Definition: nanoflann.hpp:1360
PooledAllocator pool_
Definition: nanoflann.hpp:992
NodePtr divideTree(Derived &obj, const Offset left, const Offset right, BoundingBox &bbox)
Definition: nanoflann.hpp:1039
ElementType dataset_get(const Derived &obj, IndexType element, Dimension component) const
Helper accessor to the dataset points:
Definition: nanoflann.hpp:1001
Size veclen(const Derived &obj)
Definition: nanoflann.hpp:998
Definition: nanoflann.hpp:1434
Size radiusSearchCustomCallback(const ElementType *query_point, SEARCH_CALLBACK &resultSet, const SearchParameters &searchParams={}) const
Definition: nanoflann.hpp:1675
KDTreeSingleIndexAdaptor(const KDTreeSingleIndexAdaptor< Distance, DatasetAdaptor, DIM, IndexType > &)=delete
KDTreeSingleIndexAdaptor(const Dimension dimensionality, const DatasetAdaptor &inputData, const KDTreeSingleIndexAdaptorParams &params, Args &&... args)
Definition: nanoflann.hpp:1494
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< ResultItem< IndexType, DistanceType > > &IndicesDists, const SearchParameters &searchParams={}) const
Definition: nanoflann.hpp:1654
typename Base::distance_vector_t distance_vector_t
Definition: nanoflann.hpp:1471
void loadIndex(std::istream &stream)
Definition: nanoflann.hpp:1829
typename Base::BoundingBox BoundingBox
Definition: nanoflann.hpp:1467
Size knnSearch(const ElementType *query_point, const Size num_closest, IndexType *out_indices, DistanceType *out_distances) const
Definition: nanoflann.hpp:1625
void init_vind()
Definition: nanoflann.hpp:1688
const DatasetAdaptor & dataset_
Definition: nanoflann.hpp:1442
void saveIndex(std::ostream &stream) const
Definition: nanoflann.hpp:1819
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParameters &searchParams={}) const
Definition: nanoflann.hpp:1588
void buildIndex()
Definition: nanoflann.hpp:1544
bool searchLevel(RESULTSET &result_set, const ElementType *vec, const NodePtr node, DistanceType mindist, distance_vector_t &dists, const float epsError) const
Definition: nanoflann.hpp:1736
Definition: nanoflann.hpp:1878
Size radiusSearch(const ElementType *query_point, const DistanceType &radius, std::vector< ResultItem< IndexType, DistanceType > > &IndicesDists, const SearchParameters &searchParams={}) const
Definition: nanoflann.hpp:2090
KDTreeSingleIndexDynamicAdaptor_(const Dimension dimensionality, const DatasetAdaptor &inputData, std::vector< int > &treeIndex, const KDTreeSingleIndexAdaptorParams &params=KDTreeSingleIndexAdaptorParams())
Definition: nanoflann.hpp:1930
typename Base::BoundingBox BoundingBox
Definition: nanoflann.hpp:1909
const DatasetAdaptor & dataset_
The source of our data.
Definition: nanoflann.hpp:1883
Size radiusSearchCustomCallback(const ElementType *query_point, SEARCH_CALLBACK &resultSet, const SearchParameters &searchParams={}) const
Definition: nanoflann.hpp:2111
KDTreeSingleIndexDynamicAdaptor_(const KDTreeSingleIndexDynamicAdaptor_ &rhs)=default
void buildIndex()
Definition: nanoflann.hpp:1981
void saveIndex(std::ostream &stream)
Definition: nanoflann.hpp:2236
typename Base::distance_vector_t distance_vector_t
Definition: nanoflann.hpp:1913
void loadIndex(std::istream &stream)
Definition: nanoflann.hpp:2243
Size knnSearch(const ElementType *query_point, const Size num_closest, IndexType *out_indices, DistanceType *out_distances) const
Definition: nanoflann.hpp:2061
void searchLevel(RESULTSET &result_set, const ElementType *vec, const NodePtr node, DistanceType mindist, distance_vector_t &dists, const float epsError) const
Definition: nanoflann.hpp:2161
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParameters &searchParams={}) const
Definition: nanoflann.hpp:2027
KDTreeSingleIndexDynamicAdaptor_ operator=(const KDTreeSingleIndexDynamicAdaptor_ &rhs)
Definition: nanoflann.hpp:1962
Definition: nanoflann.hpp:2264
bool findNeighbors(RESULTSET &result, const ElementType *vec, const SearchParameters &searchParams={}) const
Definition: nanoflann.hpp:2438
const DatasetAdaptor & dataset_
The source of our data.
Definition: nanoflann.hpp:2284
void removePoint(size_t idx)
Definition: nanoflann.hpp:2414
void addPoints(IndexType start, IndexType end)
Definition: nanoflann.hpp:2373
KDTreeSingleIndexDynamicAdaptor(const int dimensionality, const DatasetAdaptor &inputData, const KDTreeSingleIndexAdaptorParams &params=KDTreeSingleIndexAdaptorParams(), const size_t maximumPointCount=1000000000U)
Definition: nanoflann.hpp:2349
std::vector< int > treeIndex_
Definition: nanoflann.hpp:2288
const std::vector< index_container_t > & getAllIndices() const
Definition: nanoflann.hpp:2302
Dimension dim_
Dimensionality of each data point.
Definition: nanoflann.hpp:2293
KDTreeSingleIndexDynamicAdaptor(const KDTreeSingleIndexDynamicAdaptor< Distance, DatasetAdaptor, DIM, IndexType > &)=delete
Definition: nanoflann.hpp:165
bool addPoint(DistanceType dist, IndexType index)
Definition: nanoflann.hpp:201
Definition: nanoflann.hpp:746
~PooledAllocator()
Definition: nanoflann.hpp:784
void free_all()
Definition: nanoflann.hpp:787
void * malloc(const size_t req_size)
Definition: nanoflann.hpp:803
T * allocate(const size_t count=1)
Definition: nanoflann.hpp:860
PooledAllocator()
Definition: nanoflann.hpp:779
Definition: nanoflann.hpp:276
ResultItem< IndexType, DistanceType > worst_item() const
Definition: nanoflann.hpp:319
bool addPoint(DistanceType dist, IndexType index)
Definition: nanoflann.hpp:307
std::enable_if< has_assign< Container >::value, void >::type assign(Container &c, const size_t nElements, const T &value)
Definition: nanoflann.hpp:143
T pi_const()
Definition: nanoflann.hpp:86
std::enable_if< has_resize< Container >::value, void >::type resize(Container &c, const size_t nElements)
Definition: nanoflann.hpp:121
Definition: nanoflann.hpp:241
bool operator()(const PairType &p1, const PairType &p2) const
Definition: nanoflann.hpp:244
Definition: nanoflann.hpp:958
Definition: nanoflann.hpp:933
DistanceType divlow
The values used for subdivision.
Definition: nanoflann.hpp:946
Node * child1
Definition: nanoflann.hpp:951
Dimension divfeat
Definition: nanoflann.hpp:944
Offset right
Indices of points in leaf node.
Definition: nanoflann.hpp:940
union nanoflann::KDTreeBaseClass::Node::@0 node_type
Definition: nanoflann.hpp:2479
void query(const num_t *query_point, const Size num_closest, IndexType *out_indices, num_t *out_distances) const
Definition: nanoflann.hpp:2537
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:2501
KDTreeEigenMatrixAdaptor(const self_t &)=delete
typename index_t::Offset Offset
Definition: nanoflann.hpp:2496
Definition: nanoflann.hpp:697
Definition: nanoflann.hpp:386
Definition: nanoflann.hpp:448
Definition: nanoflann.hpp:513
Definition: nanoflann.hpp:369
Definition: nanoflann.hpp:260
DistanceType second
Distance from sample to query point.
Definition: nanoflann.hpp:268
IndexType first
Index of the sample in the dataset.
Definition: nanoflann.hpp:267
Definition: nanoflann.hpp:558
DistanceType accum_dist(const U a, const V b, const size_t) const
Definition: nanoflann.hpp:576
Definition: nanoflann.hpp:603
Definition: nanoflann.hpp:716
bool sorted
distance (default: true)
Definition: nanoflann.hpp:723
float eps
search for eps-approximate neighbours (default: 0)
Definition: nanoflann.hpp:722
Definition: nanoflann.hpp:876
Definition: nanoflann.hpp:108
Definition: nanoflann.hpp:97
Definition: nanoflann.hpp:633
Definition: nanoflann.hpp:630
Definition: nanoflann.hpp:643
Definition: nanoflann.hpp:653
Definition: nanoflann.hpp:650
Definition: nanoflann.hpp:640
Definition: nanoflann.hpp:662
Definition: nanoflann.hpp:659
Definition: nanoflann.hpp:671
Definition: nanoflann.hpp:668