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