Use rbush to speed up probe checks

index.js

@@ -1,6 +1,7 @@
 'use strict';
 
 var Queue = require('tinyqueue');
+var rbush = require('rbush');
 
 module.exports = polylabel;
 
@@ -17,6 +18,8 @@ function polylabel(polygon, precision, debug) {
         if (!i || p[1] > maxY) maxY = p[1];
     }
 
+    var tree = indexPolygon(polygon);
+
     var width = maxX - minX;
     var height = maxY - minY;
     var cellSize = Math.min(width, height);
@@ -28,15 +31,15 @@ function polylabel(polygon, precision, debug) {
     // cover polygon with initial cells
     for (var x = minX; x < maxX; x += cellSize) {
         for (var y = minY; y < maxY; y += cellSize) {
-            cellQueue.push(new Cell(x + h, y + h, h, polygon));
+            cellQueue.push(new Cell(x + h, y + h, h, tree));
         }
     }
 
     // take centroid as the first best guess
-    var bestCell = getCentroidCell(polygon);
+    var bestCell = getCentroidCell(polygon[0], tree);
 
     // special case for rectangular polygons
-    var bboxCell = new Cell(minX + width / 2, minY + height / 2, 0, polygon);
+    var bboxCell = new Cell(minX + width / 2, minY + height / 2, 0, tree);
     if (bboxCell.d > bestCell.d) bestCell = bboxCell;
 
     var numProbes = cellQueue.length;
@@ -56,10 +59,10 @@ function polylabel(polygon, precision, debug) {
 
         // split the cell into four cells
         h = cell.h / 2;
-        cellQueue.push(new Cell(cell.x - h, cell.y - h, h, polygon));
-        cellQueue.push(new Cell(cell.x + h, cell.y - h, h, polygon));
-        cellQueue.push(new Cell(cell.x - h, cell.y + h, h, polygon));
-        cellQueue.push(new Cell(cell.x + h, cell.y + h, h, polygon));
+        cellQueue.push(new Cell(cell.x - h, cell.y - h, h, tree));
+        cellQueue.push(new Cell(cell.x + h, cell.y - h, h, tree));
+        cellQueue.push(new Cell(cell.x - h, cell.y + h, h, tree));
+        cellQueue.push(new Cell(cell.x + h, cell.y + h, h, tree));
         numProbes += 4;
     }
 
@@ -71,46 +74,115 @@ function polylabel(polygon, precision, debug) {
     return [bestCell.x, bestCell.y];
 }
 
+function indexPolygon(polygon) {
+    var edges = [];
+    for (var k = 0; k < polygon.length; k++) {
+        var ring = polygon[k];
+
+        for (var i = 0, len = ring.length, j = len - 1; i < len; j = i++) {
+            var a = ring[i];
+            var b = ring[j];
+
+            edges.push({
+                minX: Math.min(a[0], b[0]),
+                minY: Math.min(a[1], b[1]),
+                maxX: Math.max(a[0], b[0]),
+                maxY: Math.max(a[1], b[1]),
+                a: a,
+                b: b
+            });
+        }
+    }
+    return rbush().load(edges);
+}
+
 function compareMax(a, b) {
     return b.max - a.max;
 }
 
-function Cell(x, y, h, polygon) {
+function Cell(x, y, h, tree) {
     this.x = x; // cell center x
     this.y = y; // cell center y
     this.h = h; // half the cell size
-    this.d = pointToPolygonDist(x, y, polygon); // distance from cell center to polygon
+    this.d = pointToPolygonDist(x, y, tree); // distance from cell center to polygon
     this.max = this.d + this.h * Math.SQRT2; // max distance to polygon within a cell
 }
 
 // signed distance from point to polygon outline (negative if point is outside)
-function pointToPolygonDist(x, y, polygon) {
+function pointToPolygonDist(x, y, tree) {
     var inside = false;
-    var minDistSq = Infinity;
 
-    for (var k = 0; k < polygon.length; k++) {
-        var ring = polygon[k];
+    var edges = tree.search({
+        minX: x,
+        minY: y,
+        maxX: Infinity,
+        maxY: y
+    });
+
+    for (var i = 0; i < edges.length; i++) {
+        var a = edges[i].a;
+        var b = edges[i].b;
+        if ((a[1] > y !== b[1] > y) &&
+            (x < (b[0] - a[0]) * (y - a[1]) / (b[1] - a[1]) + a[0])) inside = !inside;
+    }
 
-        for (var i = 0, len = ring.length, j = len - 1; i < len; j = i++) {
-            var a = ring[i];
-            var b = ring[j];
+    var minDistSq = distToClosestEdgeSq(x, y, tree);
+
+    return (inside ? 1 : -1) * Math.sqrt(minDistSq);
+}
+
+function distToClosestEdgeSq(x, y, tree) {
+    var queue = new Queue(null, compareDist);
+    var node = tree.data;
 
-            if ((a[1] > y !== b[1] > y) &&
-                (x < (b[0] - a[0]) * (y - a[1]) / (b[1] - a[1]) + a[0])) inside = !inside;
+    // search through the segment R-tree with a depth-first search using a priority queue
+    // in the order of distance to the point
+    while (node) {
+        for (var i = 0; i < node.children.length; i++) {
+            var child = node.children[i];
 
-            minDistSq = Math.min(minDistSq, getSegDistSq(x, y, a, b));
+            var dist = node.leaf ?
+                getSegDistSq(x, y, child.a, child.b) :
+                getBoxDistSq(x, y, child);
+
+            queue.push({
+                node: child,
+                dist: dist
+            });
+        }
+
+        while (queue.length && !queue.peek().node.children) {
+            return queue.pop().dist;
         }
+
+        node = queue.pop();
+        if (node) node = node.node;
     }
 
-    return (inside ? 1 : -1) * Math.sqrt(minDistSq);
+    throw new Error('Shit happened.');
+}
+
+function getBoxDistSq(x, y, box) {
+    var dx = axisDist(x, box.minX, box.maxX);
+    var dy = axisDist(y, box.minY, box.maxY);
+    return dx * dx + dy * dy;
+}
+
+function axisDist(k, min, max) {
+    return k < min ? min - k :
+           k <= max ? 0 :
+           k - max;
+}
+
+function compareDist(a, b) {
+    return a.dist - b.dist;
 }
 
 // get polygon centroid
-function getCentroidCell(polygon) {
+function getCentroidCell(points, tree) {
     var area = 0;
     var x = 0;
     var y = 0;
-    var points = polygon[0];
 
     for (var i = 0, len = points.length, j = len - 1; i < len; j = i++) {
         var a = points[i];
@@ -120,7 +192,7 @@ function getCentroidCell(polygon) {
         y += (a[1] + b[1]) * f;
         area += f * 3;
     }
-    return new Cell(x / area, y / area, 0, polygon);
+    return new Cell(x / area, y / area, 0, tree);
 }
 
 // get squared distance from a point to a segment

package.json

@@ -19,6 +19,7 @@
     "extends": "mourner"
   },
   "dependencies": {
+    "rbush": "^2.0.1",
     "tinyqueue": "^1.1.0"
   }
 }