--- a/ann-kd-tree.cpp Thu Jun 11 20:48:14 2009 -0700
+++ b/ann-kd-tree.cpp Fri Jun 12 18:08:50 2009 -0700
@@ -51,6 +51,30 @@
return bp::make_tuple(indices, distances);
}
+bp::tuple k_fixed_radius_search(ANNkd_tree& kdtree, bp::list q, double sqRad, int k, double eps, bool priority = false)
+{
+ ANNpoint annq = annAllocPt(bp::len(q));
+ for (unsigned c = 0; c < bp::len(q); ++c)
+ annq[c] = bp::extract<ANNcoord>(q[c]);
+
+ ANNidxArray nn_idx = new ANNidx[k];
+ ANNdistArray dists = new ANNdist[k];
+
+ int kball = kdtree.annkFRSearch(annq, k, sqRad, nn_idx, dists, eps);
+
+ bp::list indices, distances;
+ for (unsigned i = 0; i < k; ++i)
+ {
+ if (nn_idx[i] != ANN_NULL_IDX)
+ indices.append(nn_idx[i]);
+ if (dists[i] != ANN_DIST_INF)
+ distances.append(dists[i]);
+ }
+
+ delete nn_idx; delete dists;
+ return bp::make_tuple(indices, distances, kball);
+}
+
bp::tuple ksearch(ANNkd_tree& kdtree, bp::list q, int k, double eps)
{
return search(kdtree, q, k, eps, false);
@@ -67,8 +91,9 @@
bp::class_<ANNkd_tree>("KDTree")
.def("__init__", bp::make_constructor(&init_from_list))
- .def("ksearch", &ksearch)
- .def("k_priority_search", &k_priority_search)
+ .def("kSearch", &ksearch)
+ .def("kPriSearch", &k_priority_search)
+ .def("kFRSearch", &k_fixed_radius_search)
.def("__len__", &ANNkd_tree::nPoints)
.def("dim", &ANNkd_tree::theDim)