Added k_priority_search and max_pts_visit
authorDmitriy Morozov <dmitriy@mrzv.org>
Thu, 11 Jun 2009 20:48:14 -0700
changeset 1 d6060e9b11c4
parent 0 dc27a515a0a3
child 2 cbe1fa5e2993
Added k_priority_search and max_pts_visit
ann-kd-tree.cpp
--- a/ann-kd-tree.cpp	Thu Jun 11 12:52:30 2009 -0700
+++ b/ann-kd-tree.cpp	Thu Jun 11 20:48:14 2009 -0700
@@ -26,12 +26,8 @@
     return p;
 }
 
-bp::tuple                           ksearch(ANNkd_tree& kdtree, bp::list q, int k, double eps)
+bp::tuple                           search(ANNkd_tree& kdtree, bp::list q, int k, double eps, bool priority = false)
 {
-    ANNpointArray pts = kdtree.thePoints();
-    // for (unsigned i = 0; i < kdtree.nPoints(); ++i)
-    //     std::cout << pts[i][0] << " " << pts[i][1] << " " << pts[i][2] << std::endl;
-
     ANNpoint annq = annAllocPt(bp::len(q));
     for (unsigned c = 0; c < bp::len(q); ++c)
         annq[c] = bp::extract<ANNcoord>(q[c]);
@@ -39,7 +35,10 @@
     ANNidxArray     nn_idx  = new ANNidx[k];
     ANNdistArray    dists   = new ANNdist[k];
 
-    kdtree.annkSearch(annq, k, nn_idx, dists, eps);
+    if (priority)
+        kdtree.annkPriSearch(annq, k, nn_idx, dists, eps);
+    else
+        kdtree.annkSearch(annq, k, nn_idx, dists, eps);
 
     bp::list indices, distances;
     for (unsigned i = 0; i < k; ++i)
@@ -52,6 +51,16 @@
     return bp::make_tuple(indices, distances);
 }
 
+bp::tuple                           ksearch(ANNkd_tree& kdtree, bp::list q, int k, double eps)
+{
+    return search(kdtree, q, k, eps, false);
+}
+
+bp::tuple                           k_priority_search(ANNkd_tree& kdtree, bp::list q, int k, double eps)
+{
+    return search(kdtree, q, k, eps, true);
+}
+
 
 void export_ann_kd_tree()
 {
@@ -59,7 +68,11 @@
         .def("__init__",            bp::make_constructor(&init_from_list))
 
         .def("ksearch",             &ksearch)
+        .def("k_priority_search",   &k_priority_search)
+
         .def("__len__",             &ANNkd_tree::nPoints)
         .def("dim",                 &ANNkd_tree::theDim)
     ;
+
+    bp::def("max_pts_visit",        &annMaxPtsVisit);
 }