/** \file * Provides a Python interface for the libkdtree++. * * \author Willi Richert * * * This defines a proxy to a (int, int) -> long long KD-Tree. The long * long is needed to save a reference to Python's object id(). Thereby, * you can associate Python objects with 2D integer points. * * If you want to customize it you can adapt the following: * * * Dimension of the KD-Tree point vector. * * DIM: number of dimensions. * * operator==() and operator<<(): adapt to the number of comparisons * * py-kdtree.i: Add or adapt all usages of PyArg_ParseTuple() to reflect the * number of dimensions. * * adapt query_records in find_nearest() and count_within_range() * * Type of points. * * coord_t: If you want to have e.g. floats you have * to adapt all usages of PyArg_ParseTuple(): Change "i" to "f" e.g. * * Type of associated data. * * data_t: currently unsigned long long, which is "L" in py-kdtree.i * * PyArg_ParseTuple() has to be changed to reflect changes in data_t * */ #ifndef _PY_KDTREE_H_ #define _PY_KDTREE_H_ #include #include #include #include template struct record_t { static const size_t dim = DIM; typedef COORD_T coord_t; typedef DATA_T data_t; typedef coord_t point_t[dim]; inline coord_t operator[](size_t const N) const { return point[N]; } point_t point; data_t data; }; typedef double RANGE_T; %%TMPL_HPP_DEFS%% //////////////////////////////////////////////////////////////////////////////// // END OF TYPE SPECIFIC DEFINITIONS //////////////////////////////////////////////////////////////////////////////// template inline double tac(RECORD_T r, int k) { return r[k]; } template class PyKDTree { public: typedef record_t RECORD_T; typedef KDTree::KDTree > TREE_T; TREE_T tree; PyKDTree() : tree(std::ptr_fun(tac)) { }; void add(RECORD_T T) { tree.insert(T); }; /** Exact erase. */ bool remove(RECORD_T T) { bool removed = false; typename TREE_T::const_iterator it = tree.find_exact(T); if (it!=tree.end()) { tree.erase_exact(T); removed = true; } return removed; }; int size(void) { return tree.size(); } void optimize(void) { tree.optimise(); } RECORD_T* find_exact(RECORD_T T) { RECORD_T* found = NULL; typename TREE_T::const_iterator it = tree.find_exact(T); if (it!=tree.end()) found = new RECORD_T(*it); return found; } size_t count_within_range(typename RECORD_T::point_t T, RANGE_T range) { RECORD_T query_record; memcpy(query_record.point, T, sizeof(COORD_T)*DIM); return tree.count_within_range(query_record, range); } std::vector* find_within_range(typename RECORD_T::point_t T, RANGE_T range) { RECORD_T query_record; memcpy(query_record.point, T, sizeof(COORD_T)*DIM); std::vector *v = new std::vector; tree.find_within_range(query_record, range, std::back_inserter(*v)); return v; } RECORD_T* find_nearest (typename RECORD_T::point_t T) { RECORD_T* found = NULL; RECORD_T query_record; memcpy(query_record.point, T, sizeof(COORD_T)*DIM); std::pair best = tree.find_nearest(query_record, std::numeric_limits::max()); if (best.first!=tree.end()) { found = new RECORD_T(*best.first); } return found; } std::vector* get_all() { std::vector* v = new std::vector; for (typename TREE_T::const_iterator iter=tree.begin(); iter!=tree.end(); ++iter) { v->push_back(*iter); } return v; } size_t __len__() { return tree.size(); } }; #endif //_PY_KDTREE_H_