File Explorer

1/**2 * Mnemonist KDTree3 * =================4 *5 * Low-level JavaScript implementation of a k-dimensional tree.6 */7var iterables = require('./utils/iterables.js');8var typed = require('./utils/typed-arrays.js');9var createTupleComparator = require('./utils/comparators.js').createTupleComparator;10var FixedReverseHeap = require('./fixed-reverse-heap.js');11var inplaceQuickSortIndices = require('./sort/quick.js').inplaceQuickSortIndices;12 13/**14 * Helper function used to compute the squared distance between a query point15 * and an indexed points whose values are stored in a tree's axes.16 *17 * Note that squared distance is used instead of euclidean to avoid18 * costly sqrt computations.19 *20 * @param  {number} dimensions - Number of dimensions.21 * @param  {array}  axes       - Axes data.22 * @param  {number} pivot      - Pivot.23 * @param  {array}  point      - Query point.24 * @return {number}25 */26function squaredDistanceAxes(dimensions, axes, pivot, b) {27  var d;28 29  var dist = 0,30      step;31 32  for (d = 0; d < dimensions; d++) {33    step = axes[d][pivot] - b[d];34    dist += step * step;35  }36 37  return dist;38}39 40/**41 * Helper function used to reshape input data into low-level axes data.42 *43 * @param  {number} dimensions - Number of dimensions.44 * @param  {array}  data       - Data in the shape [label, [x, y, z...]]45 * @return {object}46 */47function reshapeIntoAxes(dimensions, data) {48  var l = data.length;49 50  var axes = new Array(dimensions),51      labels = new Array(l),52      axis;53 54  var PointerArray = typed.getPointerArray(l);55 56  var ids = new PointerArray(l);57 58  var d, i, row;59 60  var f = true;61 62  for (d = 0; d < dimensions; d++) {63    axis = new Float64Array(l);64 65    for (i = 0; i < l; i++) {66      row = data[i];67      axis[i] = row[1][d];68 69      if (f) {70        labels[i] = row[0];71        ids[i] = i;72      }73    }74 75    f = false;76    axes[d] = axis;77  }78 79  return {axes: axes, ids: ids, labels: labels};80}81 82/**83 * Helper function used to build a kd-tree from axes data.84 *85 * @param  {number} dimensions - Number of dimensions.86 * @param  {array}  axes       - Axes.87 * @param  {array}  ids        - Indices to sort.88 * @param  {array}  labels     - Point labels.89 * @return {object}90 */91function buildTree(dimensions, axes, ids, labels) {92  var l = labels.length;93 94  // NOTE: +1 because we need to keep 0 as null pointer95  var PointerArray = typed.getPointerArray(l + 1);96 97  // Building the tree98  var pivots = new PointerArray(l),99      lefts = new PointerArray(l),100      rights = new PointerArray(l);101 102  var stack = [[0, 0, ids.length, -1, 0]],103      step,104      parent,105      direction,106      median,107      pivot,108      lo,109      hi;110 111  var d, i = 0;112 113  while (stack.length !== 0) {114    step = stack.pop();115 116    d = step[0];117    lo = step[1];118    hi = step[2];119    parent = step[3];120    direction = step[4];121 122    inplaceQuickSortIndices(axes[d], ids, lo, hi);123 124    l = hi - lo;125    median = lo + (l >>> 1); // Fancy floor(l / 2)126    pivot = ids[median];127    pivots[i] = pivot;128 129    if (parent > -1) {130      if (direction === 0)131        lefts[parent] = i + 1;132      else133        rights[parent] = i + 1;134    }135 136    d = (d + 1) % dimensions;137 138    // Right139    if (median !== lo && median !== hi - 1) {140      stack.push([d, median + 1, hi, i, 1]);141    }142 143    // Left144    if (median !== lo) {145      stack.push([d, lo, median, i, 0]);146    }147 148    i++;149  }150 151  return {152    axes: axes,153    labels: labels,154    pivots: pivots,155    lefts: lefts,156    rights: rights157  };158}159 160/**161 * KDTree.162 *163 * @constructor164 */165function KDTree(dimensions, build) {166  this.dimensions = dimensions;167  this.visited = 0;168 169  this.axes = build.axes;170  this.labels = build.labels;171 172  this.pivots = build.pivots;173  this.lefts = build.lefts;174  this.rights = build.rights;175 176  this.size = this.labels.length;177}178 179/**180 * Method returning the query's nearest neighbor.181 *182 * @param  {array}  query - Query point.183 * @return {any}184 */185KDTree.prototype.nearestNeighbor = function(query) {186  var bestDistance = Infinity,187      best = null;188 189  var dimensions = this.dimensions,190      axes = this.axes,191      pivots = this.pivots,192      lefts = this.lefts,193      rights = this.rights;194 195  var visited = 0;196 197  function recurse(d, node) {198    visited++;199 200    var left = lefts[node],201        right = rights[node],202        pivot = pivots[node];203 204    var dist = squaredDistanceAxes(205      dimensions,206      axes,207      pivot,208      query209    );210 211    if (dist < bestDistance) {212      best = pivot;213      bestDistance = dist;214 215      if (dist === 0)216        return;217    }218 219    var dx = axes[d][pivot] - query[d];220 221    d = (d + 1) % dimensions;222 223    // Going the correct way?224    if (dx > 0) {225      if (left !== 0)226        recurse(d, left - 1);227    }228    else {229      if (right !== 0)230        recurse(d, right - 1);231    }232 233    // Going the other way?234    if (dx * dx < bestDistance) {235      if (dx > 0) {236        if (right !== 0)237          recurse(d, right - 1);238      }239      else {240        if (left !== 0)241          recurse(d, left - 1);242      }243    }244  }245 246  recurse(0, 0);247 248  this.visited = visited;249  return this.labels[best];250};251 252var KNN_HEAP_COMPARATOR_3 = createTupleComparator(3);253var KNN_HEAP_COMPARATOR_2 = createTupleComparator(2);254 255/**256 * Method returning the query's k nearest neighbors.257 *258 * @param  {number} k     - Number of nearest neighbor to retrieve.259 * @param  {array}  query - Query point.260 * @return {array}261 */262 263// TODO: can do better by improving upon static-kdtree here264KDTree.prototype.kNearestNeighbors = function(k, query) {265  if (k <= 0)266    throw new Error('mnemonist/kd-tree.kNearestNeighbors: k should be a positive number.');267 268  k = Math.min(k, this.size);269 270  if (k === 1)271    return [this.nearestNeighbor(query)];272 273  var heap = new FixedReverseHeap(Array, KNN_HEAP_COMPARATOR_3, k);274 275  var dimensions = this.dimensions,276      axes = this.axes,277      pivots = this.pivots,278      lefts = this.lefts,279      rights = this.rights;280 281  var visited = 0;282 283  function recurse(d, node) {284    var left = lefts[node],285        right = rights[node],286        pivot = pivots[node];287 288    var dist = squaredDistanceAxes(289      dimensions,290      axes,291      pivot,292      query293    );294 295    heap.push([dist, visited++, pivot]);296 297    var point = query[d],298        split = axes[d][pivot],299        dx = point - split;300 301    d = (d + 1) % dimensions;302 303    // Going the correct way?304    if (point < split) {305      if (left !== 0) {306        recurse(d, left - 1);307      }308    }309    else {310      if (right !== 0) {311        recurse(d, right - 1);312      }313    }314 315    // Going the other way?316    if (dx * dx < heap.peek()[0] || heap.size < k) {317      if (point < split) {318        if (right !== 0) {319          recurse(d, right - 1);320        }321      }322      else {323        if (left !== 0) {324          recurse(d, left - 1);325        }326      }327    }328  }329 330  recurse(0, 0);331 332  this.visited = visited;333 334  var best = heap.consume();335 336  for (var i = 0; i < best.length; i++)337    best[i] = this.labels[best[i][2]];338 339  return best;340};341 342/**343 * Method returning the query's k nearest neighbors by linear search.344 *345 * @param  {number} k     - Number of nearest neighbor to retrieve.346 * @param  {array}  query - Query point.347 * @return {array}348 */349KDTree.prototype.linearKNearestNeighbors = function(k, query) {350  if (k <= 0)351    throw new Error('mnemonist/kd-tree.kNearestNeighbors: k should be a positive number.');352 353  k = Math.min(k, this.size);354 355  var heap = new FixedReverseHeap(Array, KNN_HEAP_COMPARATOR_2, k);356 357  var i, l, dist;358 359  for (i = 0, l = this.size; i < l; i++) {360    dist = squaredDistanceAxes(361      this.dimensions,362      this.axes,363      this.pivots[i],364      query365    );366 367    heap.push([dist, i]);368  }369 370  var best = heap.consume();371 372  for (i = 0; i < best.length; i++)373    best[i] = this.labels[this.pivots[best[i][1]]];374 375  return best;376};377 378/**379 * Convenience known methods.380 */381KDTree.prototype.inspect = function() {382  var dummy = new Map();383 384  dummy.dimensions = this.dimensions;385 386  Object.defineProperty(dummy, 'constructor', {387    value: KDTree,388    enumerable: false389  });390 391  var i, j, point;392 393  for (i = 0; i < this.size; i++) {394    point = new Array(this.dimensions);395 396    for (j = 0; j < this.dimensions; j++)397      point[j] = this.axes[j][i];398 399    dummy.set(this.labels[i], point);400  }401 402  return dummy;403};404 405if (typeof Symbol !== 'undefined')406  KDTree.prototype[Symbol.for('nodejs.util.inspect.custom')] = KDTree.prototype.inspect;407 408/**409 * Static @.from function taking an arbitrary iterable & converting it into410 * a structure.411 *412 * @param  {Iterable} iterable   - Target iterable.413 * @param  {number}   dimensions - Space dimensions.414 * @return {KDTree}415 */416KDTree.from = function(iterable, dimensions) {417  var data = iterables.toArray(iterable);418 419  var reshaped = reshapeIntoAxes(dimensions, data);420 421  var result = buildTree(dimensions, reshaped.axes, reshaped.ids, reshaped.labels);422 423  return new KDTree(dimensions, result);424};425 426/**427 * Static @.from function building a KDTree from given axes.428 *429 * @param  {Iterable} iterable   - Target iterable.430 * @param  {number}   dimensions - Space dimensions.431 * @return {KDTree}432 */433KDTree.fromAxes = function(axes, labels) {434  if (!labels)435    labels = typed.indices(axes[0].length);436 437  var dimensions = axes.length;438 439  var result = buildTree(axes.length, axes, typed.indices(labels.length), labels);440 441  return new KDTree(dimensions, result);442};443 444/**445 * Exporting.446 */447module.exports = KDTree;448