Added fixed radius search, and renamed methods to match C++ names
authorDmitriy Morozov <dmitriy@mrzv.org>
Fri, 12 Jun 2009 18:08:50 -0700
changeset 2 cbe1fa5e2993
parent 1 d6060e9b11c4
child 3 242b36ad34e8
Added fixed radius search, and renamed methods to match C++ names
ann-kd-tree.cpp
--- 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)