KDTree WrapperHEADmain

1 files changed, 169 insertions, 0 deletions

diff --git a/KDTree.h b/KDTree.h
new file mode 100644
index 0000000..d2488e7
--- /dev/null
+++ b/KDTree.h
@@ -0,0 +1,169 @@
+/*
+    A simple C++ wrapper for John Tsiombikas's kd-tree library.
+    
+    KDTree<PointT, LabelT, Dims> kdtree;
+    PointT: 
+    - must contain a constructor to recieve points via a double buffer [double*]
+    - must have a .data() method that produces a double buffer
+
+    LabelT:
+    -  modifiable
+
+    Dims:
+    - Number of dimensions in PointT
+
+    Search methods return ResultSet similar to an iterator with more contraints
+*/
+
+#pragma once
+
+#include <memory>
+
+namespace ldrs {
+    
+    #include <kdtree.h>
+    
+    template <typename _PointT, typename _LabelT, int _Dims>
+    class KDTree {
+
+        public:
+        
+        struct Pair {
+            _PointT point;
+            _LabelT label;
+
+            typedef std::shared_ptr<Pair> Ptr;
+            typedef std::shared_ptr<Pair const> ConstPtr;
+        };
+
+        private:
+
+        struct kdtree * _tree;
+        size_t _size;
+
+        static void _delete_node_callback(void * node) {
+            delete reinterpret_cast<_LabelT*>(node);
+        }
+
+        static Pair toPair(struct kdres const * res) {
+            double pos[_Dims];
+            void * data = kd_res_item(const_cast<struct kdres *>(res), pos);
+
+            Pair pl {
+                _PointT(reinterpret_cast<double *>(pos)),
+                *reinterpret_cast<_LabelT *>(data)
+            };
+
+            return pl;
+        }
+
+        public:
+
+        typedef std::shared_ptr<KDTree> Ptr;
+        typedef std::shared_ptr<KDTree const> ConstPtr;
+
+        class ResultSet {
+            struct kdres * _res;
+        public:
+            ResultSet(struct kdres * res)
+                : _res(res)
+            {}
+
+            ~ResultSet() {
+                if(_res)
+                    kd_res_free(_res);
+            }
+
+            inline bool next() { return kd_res_next(_res) != 0; }
+
+            inline bool end() const { return kd_res_end(_res) != 0; }
+
+            inline void rewind() { kd_res_rewind(_res); }
+            
+            inline int size() const { return kd_res_size(_res); }
+
+            ResultSet(const ResultSet &) = delete;
+            ResultSet& operator=(const ResultSet &) = delete;
+            ResultSet(ResultSet && rs) {
+                _res = rs._res;
+                rs._res = nullptr;
+            }
+
+            ResultSet& operator=(ResultSet && rs) {
+                _res = rs._res;
+                rs._res = nullptr;
+                return *this;
+            }
+
+            ResultSet& operator++() {
+                next();
+                return *this;
+            }
+
+            _PointT point() const {
+                double pos[_Dims];
+                kd_res_item(_res, pos);
+                return _PointT(pos);
+            }
+
+            _LabelT& label() const {
+                return *reinterpret_cast<_LabelT*>(kd_res_item_data(_res));
+            }
+
+            Pair pair() const {
+                return KDTree::toPair(_res);
+            }
+        };
+
+        KDTree()
+            : _tree(kd_create(_Dims))
+            , _size(0)
+        {
+            kd_data_destructor(_tree, &_delete_node_callback);
+        }
+
+        ~KDTree() {
+            kd_free(_tree);
+        }
+
+        size_t size() const {
+            return _size;
+        }
+
+        void clear() {
+            _size = 0;
+            kd_clear(_tree);
+        }
+
+        int insert(_PointT const & pt, _LabelT const & label) {
+            _size++;
+
+            return kd_insert(
+                _tree,
+                pt.data(), 
+                reinterpret_cast<void *>(new _LabelT(label))
+            );
+        }
+
+        ResultSet nearestRange(_PointT const & pt, double range) const {
+            kdres * res = kd_nearest_range(_tree, pt.data(), range);
+            return ResultSet(res);
+        }
+
+        ResultSet nearestPoint(_PointT const & pt) const {
+            kdres * res = kd_nearest(_tree, pt.data());
+            return ResultSet(res);
+        }
+
+        Pair nearestExistentPoint(_PointT const & pt) const {
+            kdres * res = kd_nearest(_tree, pt.data());
+
+            Pair pl = KDTree::toPair(res);
+            
+            kd_res_free(res);
+
+            return pl;
+        }
+
+    };
+};
\ No newline at end of file