echarts KDTree 源码

echarts KDTree 代码

文件路径：/src/util/KDTree.ts

/*
* 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.
*/

import quickSelect from './quickSelect';
import { VectorArray } from 'zrender/src/core/vector';

type KDTreePoint = {
    array: VectorArray
};

class KDTreeNode<T> {

    left: KDTreeNode<T>;
    right: KDTreeNode<T>;

    axis: number;
    data: T;

    constructor(axis: number, data: T) {
        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
 */

class KDTree<T extends KDTreePoint> {

    dimension: number;
    root: KDTreeNode<T>;

    // Use one stack to avoid allocation
    // each time searching the nearest point
    private _stack: KDTreeNode<T>[] = [];
    // Again avoid allocating a new array
    // each time searching nearest N points
    private _nearstNList: {
        dist: number,
        node: KDTreeNode<T>
    }[] = [];

    constructor(points: T[], dimension?: number) {
        if (!points.length) {
            return;
        }

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

    /**
     * Resursively build the tree
     */
    private _buildTree(points: T[], left: number, right: number, axis: number): KDTreeNode<T> {
        if (right < left) {
            return null;
        }

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

        const node = new KDTreeNode(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  target Target point
     * @param  squaredDistance Squared distance function
     * @return Nearest point
     */
    nearest(target: T, squaredDistance: (a: T, b: T) => number) {
        let curr = this.root;
        const stack = this._stack;
        let idx = 0;
        let minDist = Infinity;
        let 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];
            let currDist = target.array[curr.axis] - curr.data.array[curr.axis];
            const isLeft = currDist < 0;
            let 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;
    };

    _addNearest(found: number, dist: number, node: KDTreeNode<T>) {
        const nearestNList = this._nearstNList;
        let i = found - 1;
        // Insert to the right position
        // Sort from small to large
        for (; 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  target Target point
     * @param  N
     * @param  squaredDistance Squared distance function
     * @param  output Output nearest N points
     */
    nearestN(
        target: T,
        N: number,
        squaredDistance: (a: T, b: T) => number,
        output: T[]
    ) {
        if (N <= 0) {
            output.length = 0;
            return output;
        }

        let curr = this.root;
        const stack = this._stack;
        let idx = 0;

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

        let 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];
            let currDist = target.array[curr.axis] - curr.data.array[curr.axis];
            const isLeft = currDist < 0;
            let 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 (let i = 0; i < found; i++) {
            output[i] = nearestNList[i].node.data;
        }
        output.length = found;

        return output;
    }

}


export default KDTree;

相关信息

echarts 源码目录

相关文章

echarts ECEventProcessor 源码

echarts animation 源码

echarts clazz 源码

echarts component 源码

echarts conditionalExpression 源码

echarts decal 源码

echarts event 源码

echarts format 源码

echarts graphic 源码

echarts innerStore 源码