Improve kd-tree implementation, and fix tester

Author: SuprDewd

TODO.md

@@ -2,7 +2,6 @@ Things to do
 ============
 
 * Fix convex hull edge case (some small edge case that included at least three collinear lines would return the same point on both the upper and lower hull)
-* Fix 2d-Tree (for some reason it's segfaulting on tests)
 * Shortest Path Faster Algorithm
 * Prim's algorithm
 * More formulas for geometry, trigonometry, сombinatorics, number theory

code/data-structures/kd_tree.cpp

@@ -65,31 +65,24 @@ template <int K> struct kd_tree {
   void clear() { _clr(root); root = NULL; }
   void _clr(node *n) {
     if (n) _clr(n->l), _clr(n->r), delete n; }
-  pt nearest_neighbour(const pt &p, bool allow_same=true) {
-    assert(root);
+  pair<pt, bool> nearest_neighbour(const pt &p,
+      bool allow_same=true) {
     double mn = INFINITY, cs[K];
     rep(i,0,K) cs[i] = -INFINITY;
     pt from(cs);
     rep(i,0,K) cs[i] = INFINITY;
-    pt to(cs);
-    return _nn(p, root, bb(from, to), mn, 0, allow_same).first;
-  }
-  pair<pt, bool> _nn(const pt &p, node *n, bb b,
-      double &mn, int c, bool same) {
-    if (!n || b.dist(p) > mn) return make_pair(pt(), false);
-    bool found = same || p.dist(n->p) > EPS,
-         l1 = true, l2 = false;
-    pt resp = n->p;
-    if (found) mn = min(mn, p.dist(resp));
+    pt to(cs), resp;
+    _nn(p, root, bb(from, to), mn, resp, 0, allow_same);
+    return make_pair(resp, !std::isinf(mn)); }
+  void _nn(const pt &p, node *n, bb b,
+      double &mn, pt &resp, int c, bool same) {
+    if (!n || b.dist(p) > mn) return;
+    bool l1 = true, l2 = false;
+    if ((same || p.dist(n->p) > EPS) && p.dist(n->p) < mn)
+        mn = p.dist(resp = n->p);
     node *n1 = n->l, *n2 = n->r;
     rep(i,0,2) {
-      if (i == 1 || cmp(c)(n->p, p))
-        swap(n1, n2), swap(l1, l2);
-      pair<pt, bool> res =_nn(p, n1,
-          b.bound(n->p.coord[c], c, l1), mn, INC(c), same);
-      if (res.second &&
-          (!found || p.dist(res.first) < p.dist(resp)))
-        resp = res.first, found = true;
-    }
-    return make_pair(resp, found); } };
+      if (i == 1 || cmp(c)(n->p, p)) swap(n1,n2),swap(l1,l2);
+      _nn(p, n1, b.bound(n->p.coord[c], c, l1), mn,
+          resp, INC(c), same); } } };
 // vim: cc=60 ts=2 sts=2 sw=2:

code/data-structures/kd_tree.test.cpp

@@ -41,19 +41,20 @@ struct naive_kd_tree {
     pts.insert(p);
   }
 
-  point nearest_neighbour(const point &p, bool allow_same = true) {
+  pair<point, bool> nearest_neighbour(const point &p, bool allow_same = true) {
     double mn = INFINITY;
     point res;
+    bool found = false;
     for (typename set<point>::const_iterator it = pts.begin(); it != pts.end(); ++it) {
       double dist = p.dist_to(*it);
       if (abs(dist) < EPS && !allow_same) continue;
       if (dist < mn) {
         mn = dist;
         res = *it;
+        found = true;
       }
     }
-
-    return res;
+    return make_pair(res, found);
   }
 };
 
@@ -149,19 +150,29 @@ void test() {
         bool allow_same = rand() % 2 == 0;
         // bool allow_same = true;
 #if TREE1
-        kd_tree<CURK>::pt a = tree1.nearest_neighbour(kd_tree<CURK>::pt(pt), allow_same);
+        pair<kd_tree<CURK>::pt, bool> a = tree1.nearest_neighbour(kd_tree<CURK>::pt(pt), allow_same);
+#endif
+
+#if TREE2
+        pair<naive_kd_tree<CURK>::point, bool> b = tree2.nearest_neighbour(naive_kd_tree<CURK>::point(pt), allow_same);
 #endif
 
+#if TREE1 && TREE2
+        assert(a.second == b.second);
+#endif
+#if TREE1
+        if (!a.second) continue;
+#endif
 #if TREE2
-        naive_kd_tree<CURK>::point b = tree2.nearest_neighbour(naive_kd_tree<CURK>::point(pt), allow_same);
+        if (!b.second) continue;
 #endif
 
 #if TREE1
-        double x = a.dist(kd_tree<CURK>::pt(pt));
+        double x = a.first.dist(kd_tree<CURK>::pt(pt));
 #endif
 
 #if TREE2
-        double y = b.dist_to(naive_kd_tree<CURK>::point(pt));
+        double y = b.first.dist_to(naive_kd_tree<CURK>::point(pt));
 #endif
 
 #if TREE1 && TREE2

comprog.pdf

@@ -183,7 +183,6 @@ \section{Data Structures}
         A $k$-dimensional tree supporting fast construction, adding points, and
         nearest neighbor queries.
         \fi
-        NOTE: Not completely stable, occasionally segfaults.
         \code{data-structures/kd_tree.cpp}
     \fi