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