math/SpacialIndex.js

import KDBush from 'kdbush';
import Queue from 'tinyqueue';

/**
 * 2D spatial indexing for points using a KD-tree
 * @see https://www.npmjs.com/package/kdbush
 */
class SpacialIndex {
    constructor(getX, getY) {
        this.items = [];

        this.getX = getX;
        this.getY = getY;

        this.dirty = true;
        this.tree = null;
    }

    _update() {
        let points = this.items.map(i => {
            return [this.getX(i), this.getY(i)]
        });

        this.tree = new KDBush(points);
        this.dirty = false;
    }

    _markDirty() {
        this.dirty = true;
    }

    /**
     * Add an item
     * @param {Object} item 
     */
    add(item) {
        this.items.push(item);
        this._markDirty();
    }

    /**
     * import a new array of items
     * @param {Array} items
     */
    import(items) {
        this.items = items;
        this._markDirty();
    }

    /**
     * Remove an item
     * @param {Object} item 
     */
    remove(item) {
        let i = this.items.indexOf(item);
        if (i !== -1) {
            this.items.splice(i, 1);
            this._markDirty();
        }
    }

    /**
     * Get all points in the given rectangle
     * @param {Number} minX 
     * @param {Number} minY 
     * @param {Number} maxX 
     * @param {Number} maxY 
     * @returns {Array} array of items in the rectangle
     */
    getInRect(minX, minY, maxX, maxY) {
        if (this.dirty) this._update();
        return this.tree.range(minX, minY, maxX, maxY).map(i => this.items[i]);
    }

    /**
     * Get all points in the radius around the coordinates
     * @param {Number} x coordinate
     * @param {Number} y coordinate
     * @param {Number} radius radius
     * @returns {Array} array of items in the radius
     */
    getInRadius(x, y, radius) {
        if (this.dirty) this._update();
        return this.tree.within(x, y, radius).map(i => this.items[i]);
    }

    /**
     * Get the nearest point(s) (k-nearest-neighbour)
     * @param {Number} x coordinate
     * @param {Number} y coordinate
     * @param {Number} n amount of items to return
     * @returns {*|Array} an item if n = 1, otherwise an array
     */
    getNearest(x, y, n = 1, filter) {
        var queue = new Queue(undefined, (a, b) => a.dist - b.dist);

        let result = [];

        this.items.forEach(i => {
            let dx = this.getX(i) - x;
            let dy = this.getY(i) - y;

            queue.push({
                item: i,
                dist: dx * dx + dy * dy
            });
        });

        while (queue.length) {
            let candidate = queue.pop().item;
            if (!filter || filter(candidate)) result.push(candidate);
            if (n && result.length === n) return n === 1 ? result[0] : result;
        }

        return null;
    }
}

export default SpacialIndex