KDTree.js 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282
  1. /*
  2. * Licensed to the Apache Software Foundation (ASF) under one
  3. * or more contributor license agreements. See the NOTICE file
  4. * distributed with this work for additional information
  5. * regarding copyright ownership. The ASF licenses this file
  6. * to you under the Apache License, Version 2.0 (the
  7. * "License"); you may not use this file except in compliance
  8. * with the License. You may obtain a copy of the License at
  9. *
  10. * http://www.apache.org/licenses/LICENSE-2.0
  11. *
  12. * Unless required by applicable law or agreed to in writing,
  13. * software distributed under the License is distributed on an
  14. * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15. * KIND, either express or implied. See the License for the
  16. * specific language governing permissions and limitations
  17. * under the License.
  18. */
  19. var quickSelect = require("./quickSelect");
  20. /*
  21. * Licensed to the Apache Software Foundation (ASF) under one
  22. * or more contributor license agreements. See the NOTICE file
  23. * distributed with this work for additional information
  24. * regarding copyright ownership. The ASF licenses this file
  25. * to you under the Apache License, Version 2.0 (the
  26. * "License"); you may not use this file except in compliance
  27. * with the License. You may obtain a copy of the License at
  28. *
  29. * http://www.apache.org/licenses/LICENSE-2.0
  30. *
  31. * Unless required by applicable law or agreed to in writing,
  32. * software distributed under the License is distributed on an
  33. * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  34. * KIND, either express or implied. See the License for the
  35. * specific language governing permissions and limitations
  36. * under the License.
  37. */
  38. function Node(axis, data) {
  39. this.left = null;
  40. this.right = null;
  41. this.axis = axis;
  42. this.data = data;
  43. }
  44. /**
  45. * @constructor
  46. * @alias module:echarts/data/KDTree
  47. * @param {Array} points List of points.
  48. * each point needs an array property to repesent the actual data
  49. * @param {Number} [dimension]
  50. * Point dimension.
  51. * Default will use the first point's length as dimensiont
  52. */
  53. var KDTree = function (points, dimension) {
  54. if (!points.length) {
  55. return;
  56. }
  57. if (!dimension) {
  58. dimension = points[0].array.length;
  59. }
  60. this.dimension = dimension;
  61. this.root = this._buildTree(points, 0, points.length - 1, 0); // Use one stack to avoid allocation
  62. // each time searching the nearest point
  63. this._stack = []; // Again avoid allocating a new array
  64. // each time searching nearest N points
  65. this._nearstNList = [];
  66. };
  67. /**
  68. * Resursively build the tree
  69. */
  70. KDTree.prototype._buildTree = function (points, left, right, axis) {
  71. if (right < left) {
  72. return null;
  73. }
  74. var medianIndex = Math.floor((left + right) / 2);
  75. medianIndex = quickSelect(points, left, right, medianIndex, function (a, b) {
  76. return a.array[axis] - b.array[axis];
  77. });
  78. var median = points[medianIndex];
  79. var node = new Node(axis, median);
  80. axis = (axis + 1) % this.dimension;
  81. if (right > left) {
  82. node.left = this._buildTree(points, left, medianIndex - 1, axis);
  83. node.right = this._buildTree(points, medianIndex + 1, right, axis);
  84. }
  85. return node;
  86. };
  87. /**
  88. * Find nearest point
  89. * @param {Array} target Target point
  90. * @param {Function} squaredDistance Squared distance function
  91. * @return {Array} Nearest point
  92. */
  93. KDTree.prototype.nearest = function (target, squaredDistance) {
  94. var curr = this.root;
  95. var stack = this._stack;
  96. var idx = 0;
  97. var minDist = Infinity;
  98. var nearestNode = null;
  99. if (curr.data !== target) {
  100. minDist = squaredDistance(curr.data, target);
  101. nearestNode = curr;
  102. }
  103. if (target.array[curr.axis] < curr.data.array[curr.axis]) {
  104. // Left first
  105. curr.right && (stack[idx++] = curr.right);
  106. curr.left && (stack[idx++] = curr.left);
  107. } else {
  108. // Right first
  109. curr.left && (stack[idx++] = curr.left);
  110. curr.right && (stack[idx++] = curr.right);
  111. }
  112. while (idx--) {
  113. curr = stack[idx];
  114. var currDist = target.array[curr.axis] - curr.data.array[curr.axis];
  115. var isLeft = currDist < 0;
  116. var needsCheckOtherSide = false;
  117. currDist = currDist * currDist; // Intersecting right hyperplane with minDist hypersphere
  118. if (currDist < minDist) {
  119. currDist = squaredDistance(curr.data, target);
  120. if (currDist < minDist && curr.data !== target) {
  121. minDist = currDist;
  122. nearestNode = curr;
  123. }
  124. needsCheckOtherSide = true;
  125. }
  126. if (isLeft) {
  127. if (needsCheckOtherSide) {
  128. curr.right && (stack[idx++] = curr.right);
  129. } // Search in the left area
  130. curr.left && (stack[idx++] = curr.left);
  131. } else {
  132. if (needsCheckOtherSide) {
  133. curr.left && (stack[idx++] = curr.left);
  134. } // Search the right area
  135. curr.right && (stack[idx++] = curr.right);
  136. }
  137. }
  138. return nearestNode.data;
  139. };
  140. KDTree.prototype._addNearest = function (found, dist, node) {
  141. var nearestNList = this._nearstNList; // Insert to the right position
  142. // Sort from small to large
  143. for (var i = found - 1; i > 0; i--) {
  144. if (dist >= nearestNList[i - 1].dist) {
  145. break;
  146. } else {
  147. nearestNList[i].dist = nearestNList[i - 1].dist;
  148. nearestNList[i].node = nearestNList[i - 1].node;
  149. }
  150. }
  151. nearestNList[i].dist = dist;
  152. nearestNList[i].node = node;
  153. };
  154. /**
  155. * Find nearest N points
  156. * @param {Array} target Target point
  157. * @param {number} N
  158. * @param {Function} squaredDistance Squared distance function
  159. * @param {Array} [output] Output nearest N points
  160. */
  161. KDTree.prototype.nearestN = function (target, N, squaredDistance, output) {
  162. if (N <= 0) {
  163. output.length = 0;
  164. return output;
  165. }
  166. var curr = this.root;
  167. var stack = this._stack;
  168. var idx = 0;
  169. var nearestNList = this._nearstNList;
  170. for (var i = 0; i < N; i++) {
  171. // Allocate
  172. if (!nearestNList[i]) {
  173. nearestNList[i] = {};
  174. }
  175. nearestNList[i].dist = 0;
  176. nearestNList[i].node = null;
  177. }
  178. var currDist = squaredDistance(curr.data, target);
  179. var found = 0;
  180. if (curr.data !== target) {
  181. found++;
  182. this._addNearest(found, currDist, curr);
  183. }
  184. if (target.array[curr.axis] < curr.data.array[curr.axis]) {
  185. // Left first
  186. curr.right && (stack[idx++] = curr.right);
  187. curr.left && (stack[idx++] = curr.left);
  188. } else {
  189. // Right first
  190. curr.left && (stack[idx++] = curr.left);
  191. curr.right && (stack[idx++] = curr.right);
  192. }
  193. while (idx--) {
  194. curr = stack[idx];
  195. var currDist = target.array[curr.axis] - curr.data.array[curr.axis];
  196. var isLeft = currDist < 0;
  197. var needsCheckOtherSide = false;
  198. currDist = currDist * currDist; // Intersecting right hyperplane with minDist hypersphere
  199. if (found < N || currDist < nearestNList[found - 1].dist) {
  200. currDist = squaredDistance(curr.data, target);
  201. if ((found < N || currDist < nearestNList[found - 1].dist) && curr.data !== target) {
  202. if (found < N) {
  203. found++;
  204. }
  205. this._addNearest(found, currDist, curr);
  206. }
  207. needsCheckOtherSide = true;
  208. }
  209. if (isLeft) {
  210. if (needsCheckOtherSide) {
  211. curr.right && (stack[idx++] = curr.right);
  212. } // Search in the left area
  213. curr.left && (stack[idx++] = curr.left);
  214. } else {
  215. if (needsCheckOtherSide) {
  216. curr.left && (stack[idx++] = curr.left);
  217. } // Search the right area
  218. curr.right && (stack[idx++] = curr.right);
  219. }
  220. } // Copy to output
  221. for (var i = 0; i < found; i++) {
  222. output[i] = nearestNList[i].node.data;
  223. }
  224. output.length = found;
  225. return output;
  226. };
  227. var _default = KDTree;
  228. module.exports = _default;