--- 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);
}