You can not select more than 25 topics
			Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
		
		
		
		
		
			
		
			
				
					
					
						
							448 lines
						
					
					
						
							9.2 KiB
						
					
					
				
			
		
		
	
	
							448 lines
						
					
					
						
							9.2 KiB
						
					
					
				| /**
 | |
|  * Mnemonist KDTree
 | |
|  * =================
 | |
|  *
 | |
|  * Low-level JavaScript implementation of a k-dimensional tree.
 | |
|  */
 | |
| var iterables = require('./utils/iterables.js');
 | |
| var typed = require('./utils/typed-arrays.js');
 | |
| var createTupleComparator = require('./utils/comparators.js').createTupleComparator;
 | |
| var FixedReverseHeap = require('./fixed-reverse-heap.js');
 | |
| var inplaceQuickSortIndices = require('./sort/quick.js').inplaceQuickSortIndices;
 | |
| 
 | |
| /**
 | |
|  * Helper function used to compute the squared distance between a query point
 | |
|  * and an indexed points whose values are stored in a tree's axes.
 | |
|  *
 | |
|  * Note that squared distance is used instead of euclidean to avoid
 | |
|  * costly sqrt computations.
 | |
|  *
 | |
|  * @param  {number} dimensions - Number of dimensions.
 | |
|  * @param  {array}  axes       - Axes data.
 | |
|  * @param  {number} pivot      - Pivot.
 | |
|  * @param  {array}  point      - Query point.
 | |
|  * @return {number}
 | |
|  */
 | |
| function squaredDistanceAxes(dimensions, axes, pivot, b) {
 | |
|   var d;
 | |
| 
 | |
|   var dist = 0,
 | |
|       step;
 | |
| 
 | |
|   for (d = 0; d < dimensions; d++) {
 | |
|     step = axes[d][pivot] - b[d];
 | |
|     dist += step * step;
 | |
|   }
 | |
| 
 | |
|   return dist;
 | |
| }
 | |
| 
 | |
| /**
 | |
|  * Helper function used to reshape input data into low-level axes data.
 | |
|  *
 | |
|  * @param  {number} dimensions - Number of dimensions.
 | |
|  * @param  {array}  data       - Data in the shape [label, [x, y, z...]]
 | |
|  * @return {object}
 | |
|  */
 | |
| function reshapeIntoAxes(dimensions, data) {
 | |
|   var l = data.length;
 | |
| 
 | |
|   var axes = new Array(dimensions),
 | |
|       labels = new Array(l),
 | |
|       axis;
 | |
| 
 | |
|   var PointerArray = typed.getPointerArray(l);
 | |
| 
 | |
|   var ids = new PointerArray(l);
 | |
| 
 | |
|   var d, i, row;
 | |
| 
 | |
|   var f = true;
 | |
| 
 | |
|   for (d = 0; d < dimensions; d++) {
 | |
|     axis = new Float64Array(l);
 | |
| 
 | |
|     for (i = 0; i < l; i++) {
 | |
|       row = data[i];
 | |
|       axis[i] = row[1][d];
 | |
| 
 | |
|       if (f) {
 | |
|         labels[i] = row[0];
 | |
|         ids[i] = i;
 | |
|       }
 | |
|     }
 | |
| 
 | |
|     f = false;
 | |
|     axes[d] = axis;
 | |
|   }
 | |
| 
 | |
|   return {axes: axes, ids: ids, labels: labels};
 | |
| }
 | |
| 
 | |
| /**
 | |
|  * Helper function used to build a kd-tree from axes data.
 | |
|  *
 | |
|  * @param  {number} dimensions - Number of dimensions.
 | |
|  * @param  {array}  axes       - Axes.
 | |
|  * @param  {array}  ids        - Indices to sort.
 | |
|  * @param  {array}  labels     - Point labels.
 | |
|  * @return {object}
 | |
|  */
 | |
| function buildTree(dimensions, axes, ids, labels) {
 | |
|   var l = labels.length;
 | |
| 
 | |
|   // NOTE: +1 because we need to keep 0 as null pointer
 | |
|   var PointerArray = typed.getPointerArray(l + 1);
 | |
| 
 | |
|   // Building the tree
 | |
|   var pivots = new PointerArray(l),
 | |
|       lefts = new PointerArray(l),
 | |
|       rights = new PointerArray(l);
 | |
| 
 | |
|   var stack = [[0, 0, ids.length, -1, 0]],
 | |
|       step,
 | |
|       parent,
 | |
|       direction,
 | |
|       median,
 | |
|       pivot,
 | |
|       lo,
 | |
|       hi;
 | |
| 
 | |
|   var d, i = 0;
 | |
| 
 | |
|   while (stack.length !== 0) {
 | |
|     step = stack.pop();
 | |
| 
 | |
|     d = step[0];
 | |
|     lo = step[1];
 | |
|     hi = step[2];
 | |
|     parent = step[3];
 | |
|     direction = step[4];
 | |
| 
 | |
|     inplaceQuickSortIndices(axes[d], ids, lo, hi);
 | |
| 
 | |
|     l = hi - lo;
 | |
|     median = lo + (l >>> 1); // Fancy floor(l / 2)
 | |
|     pivot = ids[median];
 | |
|     pivots[i] = pivot;
 | |
| 
 | |
|     if (parent > -1) {
 | |
|       if (direction === 0)
 | |
|         lefts[parent] = i + 1;
 | |
|       else
 | |
|         rights[parent] = i + 1;
 | |
|     }
 | |
| 
 | |
|     d = (d + 1) % dimensions;
 | |
| 
 | |
|     // Right
 | |
|     if (median !== lo && median !== hi - 1) {
 | |
|       stack.push([d, median + 1, hi, i, 1]);
 | |
|     }
 | |
| 
 | |
|     // Left
 | |
|     if (median !== lo) {
 | |
|       stack.push([d, lo, median, i, 0]);
 | |
|     }
 | |
| 
 | |
|     i++;
 | |
|   }
 | |
| 
 | |
|   return {
 | |
|     axes: axes,
 | |
|     labels: labels,
 | |
|     pivots: pivots,
 | |
|     lefts: lefts,
 | |
|     rights: rights
 | |
|   };
 | |
| }
 | |
| 
 | |
| /**
 | |
|  * KDTree.
 | |
|  *
 | |
|  * @constructor
 | |
|  */
 | |
| function KDTree(dimensions, build) {
 | |
|   this.dimensions = dimensions;
 | |
|   this.visited = 0;
 | |
| 
 | |
|   this.axes = build.axes;
 | |
|   this.labels = build.labels;
 | |
| 
 | |
|   this.pivots = build.pivots;
 | |
|   this.lefts = build.lefts;
 | |
|   this.rights = build.rights;
 | |
| 
 | |
|   this.size = this.labels.length;
 | |
| }
 | |
| 
 | |
| /**
 | |
|  * Method returning the query's nearest neighbor.
 | |
|  *
 | |
|  * @param  {array}  query - Query point.
 | |
|  * @return {any}
 | |
|  */
 | |
| KDTree.prototype.nearestNeighbor = function(query) {
 | |
|   var bestDistance = Infinity,
 | |
|       best = null;
 | |
| 
 | |
|   var dimensions = this.dimensions,
 | |
|       axes = this.axes,
 | |
|       pivots = this.pivots,
 | |
|       lefts = this.lefts,
 | |
|       rights = this.rights;
 | |
| 
 | |
|   var visited = 0;
 | |
| 
 | |
|   function recurse(d, node) {
 | |
|     visited++;
 | |
| 
 | |
|     var left = lefts[node],
 | |
|         right = rights[node],
 | |
|         pivot = pivots[node];
 | |
| 
 | |
|     var dist = squaredDistanceAxes(
 | |
|       dimensions,
 | |
|       axes,
 | |
|       pivot,
 | |
|       query
 | |
|     );
 | |
| 
 | |
|     if (dist < bestDistance) {
 | |
|       best = pivot;
 | |
|       bestDistance = dist;
 | |
| 
 | |
|       if (dist === 0)
 | |
|         return;
 | |
|     }
 | |
| 
 | |
|     var dx = axes[d][pivot] - query[d];
 | |
| 
 | |
|     d = (d + 1) % dimensions;
 | |
| 
 | |
|     // Going the correct way?
 | |
|     if (dx > 0) {
 | |
|       if (left !== 0)
 | |
|         recurse(d, left - 1);
 | |
|     }
 | |
|     else {
 | |
|       if (right !== 0)
 | |
|         recurse(d, right - 1);
 | |
|     }
 | |
| 
 | |
|     // Going the other way?
 | |
|     if (dx * dx < bestDistance) {
 | |
|       if (dx > 0) {
 | |
|         if (right !== 0)
 | |
|           recurse(d, right - 1);
 | |
|       }
 | |
|       else {
 | |
|         if (left !== 0)
 | |
|           recurse(d, left - 1);
 | |
|       }
 | |
|     }
 | |
|   }
 | |
| 
 | |
|   recurse(0, 0);
 | |
| 
 | |
|   this.visited = visited;
 | |
|   return this.labels[best];
 | |
| };
 | |
| 
 | |
| var KNN_HEAP_COMPARATOR_3 = createTupleComparator(3);
 | |
| var KNN_HEAP_COMPARATOR_2 = createTupleComparator(2);
 | |
| 
 | |
| /**
 | |
|  * Method returning the query's k nearest neighbors.
 | |
|  *
 | |
|  * @param  {number} k     - Number of nearest neighbor to retrieve.
 | |
|  * @param  {array}  query - Query point.
 | |
|  * @return {array}
 | |
|  */
 | |
| 
 | |
| // TODO: can do better by improving upon static-kdtree here
 | |
| KDTree.prototype.kNearestNeighbors = function(k, query) {
 | |
|   if (k <= 0)
 | |
|     throw new Error('mnemonist/kd-tree.kNearestNeighbors: k should be a positive number.');
 | |
| 
 | |
|   k = Math.min(k, this.size);
 | |
| 
 | |
|   if (k === 1)
 | |
|     return [this.nearestNeighbor(query)];
 | |
| 
 | |
|   var heap = new FixedReverseHeap(Array, KNN_HEAP_COMPARATOR_3, k);
 | |
| 
 | |
|   var dimensions = this.dimensions,
 | |
|       axes = this.axes,
 | |
|       pivots = this.pivots,
 | |
|       lefts = this.lefts,
 | |
|       rights = this.rights;
 | |
| 
 | |
|   var visited = 0;
 | |
| 
 | |
|   function recurse(d, node) {
 | |
|     var left = lefts[node],
 | |
|         right = rights[node],
 | |
|         pivot = pivots[node];
 | |
| 
 | |
|     var dist = squaredDistanceAxes(
 | |
|       dimensions,
 | |
|       axes,
 | |
|       pivot,
 | |
|       query
 | |
|     );
 | |
| 
 | |
|     heap.push([dist, visited++, pivot]);
 | |
| 
 | |
|     var point = query[d],
 | |
|         split = axes[d][pivot],
 | |
|         dx = point - split;
 | |
| 
 | |
|     d = (d + 1) % dimensions;
 | |
| 
 | |
|     // Going the correct way?
 | |
|     if (point < split) {
 | |
|       if (left !== 0) {
 | |
|         recurse(d, left - 1);
 | |
|       }
 | |
|     }
 | |
|     else {
 | |
|       if (right !== 0) {
 | |
|         recurse(d, right - 1);
 | |
|       }
 | |
|     }
 | |
| 
 | |
|     // Going the other way?
 | |
|     if (dx * dx < heap.peek()[0] || heap.size < k) {
 | |
|       if (point < split) {
 | |
|         if (right !== 0) {
 | |
|           recurse(d, right - 1);
 | |
|         }
 | |
|       }
 | |
|       else {
 | |
|         if (left !== 0) {
 | |
|           recurse(d, left - 1);
 | |
|         }
 | |
|       }
 | |
|     }
 | |
|   }
 | |
| 
 | |
|   recurse(0, 0);
 | |
| 
 | |
|   this.visited = visited;
 | |
| 
 | |
|   var best = heap.consume();
 | |
| 
 | |
|   for (var i = 0; i < best.length; i++)
 | |
|     best[i] = this.labels[best[i][2]];
 | |
| 
 | |
|   return best;
 | |
| };
 | |
| 
 | |
| /**
 | |
|  * Method returning the query's k nearest neighbors by linear search.
 | |
|  *
 | |
|  * @param  {number} k     - Number of nearest neighbor to retrieve.
 | |
|  * @param  {array}  query - Query point.
 | |
|  * @return {array}
 | |
|  */
 | |
| KDTree.prototype.linearKNearestNeighbors = function(k, query) {
 | |
|   if (k <= 0)
 | |
|     throw new Error('mnemonist/kd-tree.kNearestNeighbors: k should be a positive number.');
 | |
| 
 | |
|   k = Math.min(k, this.size);
 | |
| 
 | |
|   var heap = new FixedReverseHeap(Array, KNN_HEAP_COMPARATOR_2, k);
 | |
| 
 | |
|   var i, l, dist;
 | |
| 
 | |
|   for (i = 0, l = this.size; i < l; i++) {
 | |
|     dist = squaredDistanceAxes(
 | |
|       this.dimensions,
 | |
|       this.axes,
 | |
|       this.pivots[i],
 | |
|       query
 | |
|     );
 | |
| 
 | |
|     heap.push([dist, i]);
 | |
|   }
 | |
| 
 | |
|   var best = heap.consume();
 | |
| 
 | |
|   for (i = 0; i < best.length; i++)
 | |
|     best[i] = this.labels[this.pivots[best[i][1]]];
 | |
| 
 | |
|   return best;
 | |
| };
 | |
| 
 | |
| /**
 | |
|  * Convenience known methods.
 | |
|  */
 | |
| KDTree.prototype.inspect = function() {
 | |
|   var dummy = new Map();
 | |
| 
 | |
|   dummy.dimensions = this.dimensions;
 | |
| 
 | |
|   Object.defineProperty(dummy, 'constructor', {
 | |
|     value: KDTree,
 | |
|     enumerable: false
 | |
|   });
 | |
| 
 | |
|   var i, j, point;
 | |
| 
 | |
|   for (i = 0; i < this.size; i++) {
 | |
|     point = new Array(this.dimensions);
 | |
| 
 | |
|     for (j = 0; j < this.dimensions; j++)
 | |
|       point[j] = this.axes[j][i];
 | |
| 
 | |
|     dummy.set(this.labels[i], point);
 | |
|   }
 | |
| 
 | |
|   return dummy;
 | |
| };
 | |
| 
 | |
| if (typeof Symbol !== 'undefined')
 | |
|   KDTree.prototype[Symbol.for('nodejs.util.inspect.custom')] = KDTree.prototype.inspect;
 | |
| 
 | |
| /**
 | |
|  * Static @.from function taking an arbitrary iterable & converting it into
 | |
|  * a structure.
 | |
|  *
 | |
|  * @param  {Iterable} iterable   - Target iterable.
 | |
|  * @param  {number}   dimensions - Space dimensions.
 | |
|  * @return {KDTree}
 | |
|  */
 | |
| KDTree.from = function(iterable, dimensions) {
 | |
|   var data = iterables.toArray(iterable);
 | |
| 
 | |
|   var reshaped = reshapeIntoAxes(dimensions, data);
 | |
| 
 | |
|   var result = buildTree(dimensions, reshaped.axes, reshaped.ids, reshaped.labels);
 | |
| 
 | |
|   return new KDTree(dimensions, result);
 | |
| };
 | |
| 
 | |
| /**
 | |
|  * Static @.from function building a KDTree from given axes.
 | |
|  *
 | |
|  * @param  {Iterable} iterable   - Target iterable.
 | |
|  * @param  {number}   dimensions - Space dimensions.
 | |
|  * @return {KDTree}
 | |
|  */
 | |
| KDTree.fromAxes = function(axes, labels) {
 | |
|   if (!labels)
 | |
|     labels = typed.indices(axes[0].length);
 | |
| 
 | |
|   var dimensions = axes.length;
 | |
| 
 | |
|   var result = buildTree(axes.length, axes, typed.indices(labels.length), labels);
 | |
| 
 | |
|   return new KDTree(dimensions, result);
 | |
| };
 | |
| 
 | |
| /**
 | |
|  * Exporting.
 | |
|  */
 | |
| module.exports = KDTree;
 |