look-alike/lib/kdtree.js

A simple-yet-powerful KD-tree library for NodeJS, with support for lightning-fast k-Nearest Neighbour queries. Supports normalization, weights, key and filter parameters

// Generated by CoffeeScript 1.6.3
(function() {
  var KDtree, util;

  util = require('./util');

  KDtree = (function() {
    KDtree.prototype.getRoot = function() {
      return this._tree;
    };

    function KDtree(objects, options) {
      var k, v, _base, _helper,
        _this = this;
      this.objects = objects;
      this.options = options;
      if (!(arguments.length > 0)) {
        throw new Error('Need at least 1 argument');
      }
      if (Array.isArray(this.objects)) {
        if (this.objects.some(function(x) {
          return x && x.toString() !== '[object Object]';
        })) {
          throw new Error('Expecting an array of objects as first argument');
        }
      } else {
        throw new Error('Expecting an array of objects as first argument');
      }
      this.options = this.options || {};
      if ((_base = this.options).attributes == null) {
        _base.attributes = (function() {
          var _ref, _results;
          _ref = this.objects[0];
          _results = [];
          for (k in _ref) {
            v = _ref[k];
            _results.push(k);
          }
          return _results;
        }).call(this);
      }
      if (Array.isArray(this.options.attributes)) {
        if (this.options.attributes.some(function(x) {
          return typeof x !== 'string';
        })) {
          throw new Error('Expecting an array of strings for attributes');
        }
      } else {
        throw new Error('Expecting an array of strings for attributes');
      }
      if (!this.objects.every(function(x) {
        return _this.options.attributes.every(function(k) {
          if (_this.options.key) {
            return _this.options.key(x).hasOwnProperty(k);
          } else {
            return x.hasOwnProperty(k);
          }
        });
      })) {
        throw new Error("Expecting all objects to have at least the same keys as first object or second parameter");
      }
      this.stdv = util.allStdvs(this.options.attributes, this.objects, this.options.key);
      _helper = function(objects, depth) {
        var attr, bounds, len, node, o, splits, temp;
        if (!objects.length) {
          return null;
        }
        len = _this.options.attributes.length;
        attr = _this.options.attributes[depth % len];
        objects.sort(function(a, b) {
          var _ref;
          if ((_ref = this.options) != null ? _ref.key : void 0) {
            return this.options.key(a)[attr] - this.options.key(b)[attr];
          } else {
            return a[attr] - b[attr];
          }
        });
        temp = objects;
        if (_this.options.key) {
          temp = (function() {
            var _i, _len, _results;
            _results = [];
            for (_i = 0, _len = objects.length; _i < _len; _i++) {
              o = objects[_i];
              _results.push(this.options.key(o));
            }
            return _results;
          }).call(_this);
        }
        bounds = util.medianIndex((function() {
          var _i, _len, _results;
          _results = [];
          for (_i = 0, _len = temp.length; _i < _len; _i++) {
            o = temp[_i];
            _results.push(o[attr]);
          }
          return _results;
        })());
        splits = util.getSplit(objects, bounds, _this.options.attributes, attr, _this.options.key);
        return node = {
          val: splits.identicals,
          left: _helper(splits.left, depth + 1),
          right: _helper(splits.right, depth + 1)
        };
      };
      this._tree = _helper(this.objects, 0);
    }

    KDtree.prototype.query = function(subject, options) {
      var BPQ, Q, root, _helper,
        _this = this;
      if (!this.options.attributes.every(function(k) {
        return subject.hasOwnProperty(k);
      })) {
        throw new Error("Subject does not have all keys");
      }
      if (options) {
        if (!options.k) {
          options.k = 1;
        }
        if (options.normalize !== false) {
          options.normalize = true;
        }
      } else {
        options = {
          k: 1,
          normalize: true
        };
      }
      BPQ = require('./bpq');
      Q = new BPQ(options.k);
      _helper = function(node, depth) {
        var attr, attr_dist, dist, goLeft, len, o, objectValues, _i, _len, _ref, _ref1, _ref2;
        if (!node) {
          return null;
        }
        len = _this.options.attributes.length;
        attr = _this.options.attributes[depth % len];
        objectValues = ((_ref = _this.options) != null ? _ref.key : void 0) ? (_ref1 = _this.options) != null ? _ref1.key(node.val[0]) : void 0 : node.val[0];
        if (options.normalize) {
          dist = util.distance(subject, objectValues, {
            stdv: _this.stdv,
            weights: options.weights
          });
        } else {
          dist = util.distance(subject, objectValues, {
            weights: options.weights
          });
        }
        options = options || {};
        if (options.filter) {
          _ref2 = node.val;
          for (_i = 0, _len = _ref2.length; _i < _len; _i++) {
            o = _ref2[_i];
            if (options.filter(o)) {
              Q.insert(o, dist);
            }
          }
        } else {
          Q.insert(node.val, dist);
        }
        goLeft = subject[attr] < objectValues[attr];
        if (goLeft) {
          _helper(node.left, depth + 1);
        } else {
          _helper(node.right, depth + 1);
        }
        if (options.normalize) {
          attr_dist = Math.abs(objectValues[attr] - subject[attr]) / _this.stdv[attr];
        } else {
          attr_dist = Math.abs(objectValues[attr] - subject[attr]);
        }
        if (options.weights) {
          attr_dist *= options.weights[attr];
        }
        if (options.k > Q.getSize() || attr_dist < Q.getMaxPriority()) {
          if (goLeft) {
            _helper(node.right, depth + 1);
          } else {
            _helper(node.left, depth + 1);
          }
        }
        return Q.getObjects();
      };
      root = this.getRoot();
      if (root === null) {
        return [];
      } else {
        return _helper(root, 0);
      }
    };

    return KDtree;

  })();

  module.exports = KDtree;

}).call(this);