KDTree.js

/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License.  You may obtain a copy of the License at
*
*   http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied.  See the License for the
* specific language governing permissions and limitations
* under the License.
*/

/**
 * K-Dimension Tree
 *
 * @module echarts/data/KDTree
 * @author Yi Shen(https://github.com/pissang)
 */

import quickSelect from './quickSelect';

function Node(axis, data) {
    this.left = null;
    this.right = null;
    this.axis = axis;

    this.data = data;
}

/**
 * @constructor
 * @alias module:echarts/data/KDTree
 * @param {Array} points List of points.
 * each point needs an array property to repesent the actual data
 * @param {Number} [dimension]
 *        Point dimension.
 *        Default will use the first point's length as dimensiont
 */
var KDTree = function (points, dimension) {
    if (!points.length) {
        return;
    }

    if (!dimension) {
        dimension = points[0].array.length;
    }
    this.dimension = dimension;
    this.root = this._buildTree(points, 0, points.length - 1, 0);

    // Use one stack to avoid allocation
    // each time searching the nearest point
    this._stack = [];
    // Again avoid allocating a new array
    // each time searching nearest N points
    this._nearstNList = [];
};

/**
 * Resursively build the tree
 */
KDTree.prototype._buildTree = function (points, left, right, axis) {
    if (right < left) {
        return null;
    }

    var medianIndex = Math.floor((left + right) / 2);
    medianIndex = quickSelect(
        points, left, right, medianIndex,
        function (a, b) {
            return a.array[axis] - b.array[axis];
        }
    );
    var median = points[medianIndex];

    var node = new Node(axis, median);

    axis = (axis + 1) % this.dimension;
    if (right > left) {
        node.left = this._buildTree(points, left, medianIndex - 1, axis);
        node.right = this._buildTree(points, medianIndex + 1, right, axis);
    }

    return node;
};

/**
 * Find nearest point
 * @param  {Array} target Target point
 * @param  {Function} squaredDistance Squared distance function
 * @return {Array} Nearest point
 */
KDTree.prototype.nearest = function (target, squaredDistance) {
    var curr = this.root;
    var stack = this._stack;
    var idx = 0;
    var minDist = Infinity;
    var nearestNode = null;
    if (curr.data !== target) {
        minDist = squaredDistance(curr.data, target);
        nearestNode = curr;
    }

    if (target.array[curr.axis] < curr.data.array[curr.axis]) {
        // Left first
        curr.right && (stack[idx++] = curr.right);
        curr.left && (stack[idx++] = curr.left);
    }
    else {
        // Right first
        curr.left && (stack[idx++] = curr.left);
        curr.right && (stack[idx++] = curr.right);
    }

    while (idx--) {
        curr = stack[idx];
        var currDist = target.array[curr.axis] - curr.data.array[curr.axis];
        var isLeft = currDist < 0;
        var needsCheckOtherSide = false;
        currDist = currDist * currDist;
        // Intersecting right hyperplane with minDist hypersphere
        if (currDist < minDist) {
            currDist = squaredDistance(curr.data, target);
            if (currDist < minDist && curr.data !== target) {
                minDist = currDist;
                nearestNode = curr;
            }
            needsCheckOtherSide = true;
        }
        if (isLeft) {
            if (needsCheckOtherSide) {
                curr.right && (stack[idx++] = curr.right);
            }
            // Search in the left area
            curr.left && (stack[idx++] = curr.left);
        }
        else {
            if (needsCheckOtherSide) {
                curr.left && (stack[idx++] = curr.left);
            }
            // Search the right area
            curr.right && (stack[idx++] = curr.right);
        }
    }

    return nearestNode.data;
};

KDTree.prototype._addNearest = function (found, dist, node) {
    var nearestNList = this._nearstNList;

    // Insert to the right position
    // Sort from small to large
    for (var i = found - 1; i > 0; i--) {
        if (dist >= nearestNList[i - 1].dist) {
            break;
        }
        else {
            nearestNList[i].dist = nearestNList[i - 1].dist;
            nearestNList[i].node = nearestNList[i - 1].node;
        }
    }

    nearestNList[i].dist = dist;
    nearestNList[i].node = node;
};

/**
 * Find nearest N points
 * @param  {Array} target Target point
 * @param  {number} N
 * @param  {Function} squaredDistance Squared distance function
 * @param  {Array} [output] Output nearest N points
 */
KDTree.prototype.nearestN = function (target, N, squaredDistance, output) {
    if (N <= 0) {
        output.length = 0;
        return output;
    }

    var curr = this.root;
    var stack = this._stack;
    var idx = 0;

    var nearestNList = this._nearstNList;
    for (var i = 0; i < N; i++) {
        // Allocate
        if (!nearestNList[i]) {
            nearestNList[i] = {};
        }
        nearestNList[i].dist = 0;
        nearestNList[i].node = null;
    }
    var currDist = squaredDistance(curr.data, target);

    var found = 0;
    if (curr.data !== target) {
        found++;
        this._addNearest(found, currDist, curr);
    }

    if (target.array[curr.axis] < curr.data.array[curr.axis]) {
        // Left first
        curr.right && (stack[idx++] = curr.right);
        curr.left && (stack[idx++] = curr.left);
    }
    else {
        // Right first
        curr.left && (stack[idx++] = curr.left);
        curr.right && (stack[idx++] = curr.right);
    }

    while (idx--) {
        curr = stack[idx];
        var currDist = target.array[curr.axis] - curr.data.array[curr.axis];
        var isLeft = currDist < 0;
        var needsCheckOtherSide = false;
        currDist = currDist * currDist;
        // Intersecting right hyperplane with minDist hypersphere
        if (found < N || currDist < nearestNList[found - 1].dist) {
            currDist = squaredDistance(curr.data, target);
            if (
                (found < N || currDist < nearestNList[found - 1].dist)
                && curr.data !== target
            ) {
                if (found < N) {
                    found++;
                }
                this._addNearest(found, currDist, curr);
            }
            needsCheckOtherSide = true;
        }
        if (isLeft) {
            if (needsCheckOtherSide) {
                curr.right && (stack[idx++] = curr.right);
            }
            // Search in the left area
            curr.left && (stack[idx++] = curr.left);
        }
        else {
            if (needsCheckOtherSide) {
                curr.left && (stack[idx++] = curr.left);
            }
            // Search the right area
            curr.right && (stack[idx++] = curr.right);
        }
    }

    // Copy to output
    for (var i = 0; i < found; i++) {
        output[i] = nearestNList[i].node.data;
    }
    output.length = found;

    return output;
};

export default KDTree;