KDTree.js 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257
  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. import quickSelect from './quickSelect';
  20. function Node(axis, data) {
  21. this.left = null;
  22. this.right = null;
  23. this.axis = axis;
  24. this.data = data;
  25. }
  26. /**
  27. * @constructor
  28. * @alias module:echarts/data/KDTree
  29. * @param {Array} points List of points.
  30. * each point needs an array property to repesent the actual data
  31. * @param {Number} [dimension]
  32. * Point dimension.
  33. * Default will use the first point's length as dimensiont
  34. */
  35. var KDTree = function (points, dimension) {
  36. if (!points.length) {
  37. return;
  38. }
  39. if (!dimension) {
  40. dimension = points[0].array.length;
  41. }
  42. this.dimension = dimension;
  43. this.root = this._buildTree(points, 0, points.length - 1, 0);
  44. // Use one stack to avoid allocation
  45. // each time searching the nearest point
  46. this._stack = [];
  47. // Again avoid allocating a new array
  48. // each time searching nearest N points
  49. this._nearstNList = [];
  50. };
  51. /**
  52. * Resursively build the tree
  53. */
  54. KDTree.prototype._buildTree = function (points, left, right, axis) {
  55. if (right < left) {
  56. return null;
  57. }
  58. var medianIndex = Math.floor((left + right) / 2);
  59. medianIndex = quickSelect(
  60. points, left, right, medianIndex,
  61. function (a, b) {
  62. return a.array[axis] - b.array[axis];
  63. }
  64. );
  65. var median = points[medianIndex];
  66. var node = new Node(axis, median);
  67. axis = (axis + 1) % this.dimension;
  68. if (right > left) {
  69. node.left = this._buildTree(points, left, medianIndex - 1, axis);
  70. node.right = this._buildTree(points, medianIndex + 1, right, axis);
  71. }
  72. return node;
  73. };
  74. /**
  75. * Find nearest point
  76. * @param {Array} target Target point
  77. * @param {Function} squaredDistance Squared distance function
  78. * @return {Array} Nearest point
  79. */
  80. KDTree.prototype.nearest = function (target, squaredDistance) {
  81. var curr = this.root;
  82. var stack = this._stack;
  83. var idx = 0;
  84. var minDist = Infinity;
  85. var nearestNode = null;
  86. if (curr.data !== target) {
  87. minDist = squaredDistance(curr.data, target);
  88. nearestNode = curr;
  89. }
  90. if (target.array[curr.axis] < curr.data.array[curr.axis]) {
  91. // Left first
  92. curr.right && (stack[idx++] = curr.right);
  93. curr.left && (stack[idx++] = curr.left);
  94. }
  95. else {
  96. // Right first
  97. curr.left && (stack[idx++] = curr.left);
  98. curr.right && (stack[idx++] = curr.right);
  99. }
  100. while (idx--) {
  101. curr = stack[idx];
  102. var currDist = target.array[curr.axis] - curr.data.array[curr.axis];
  103. var isLeft = currDist < 0;
  104. var needsCheckOtherSide = false;
  105. currDist = currDist * currDist;
  106. // Intersecting right hyperplane with minDist hypersphere
  107. if (currDist < minDist) {
  108. currDist = squaredDistance(curr.data, target);
  109. if (currDist < minDist && curr.data !== target) {
  110. minDist = currDist;
  111. nearestNode = curr;
  112. }
  113. needsCheckOtherSide = true;
  114. }
  115. if (isLeft) {
  116. if (needsCheckOtherSide) {
  117. curr.right && (stack[idx++] = curr.right);
  118. }
  119. // Search in the left area
  120. curr.left && (stack[idx++] = curr.left);
  121. }
  122. else {
  123. if (needsCheckOtherSide) {
  124. curr.left && (stack[idx++] = curr.left);
  125. }
  126. // Search the right area
  127. curr.right && (stack[idx++] = curr.right);
  128. }
  129. }
  130. return nearestNode.data;
  131. };
  132. KDTree.prototype._addNearest = function (found, dist, node) {
  133. var nearestNList = this._nearstNList;
  134. // Insert to the right position
  135. // Sort from small to large
  136. for (var i = found - 1; i > 0; i--) {
  137. if (dist >= nearestNList[i - 1].dist) {
  138. break;
  139. }
  140. else {
  141. nearestNList[i].dist = nearestNList[i - 1].dist;
  142. nearestNList[i].node = nearestNList[i - 1].node;
  143. }
  144. }
  145. nearestNList[i].dist = dist;
  146. nearestNList[i].node = node;
  147. };
  148. /**
  149. * Find nearest N points
  150. * @param {Array} target Target point
  151. * @param {number} N
  152. * @param {Function} squaredDistance Squared distance function
  153. * @param {Array} [output] Output nearest N points
  154. */
  155. KDTree.prototype.nearestN = function (target, N, squaredDistance, output) {
  156. if (N <= 0) {
  157. output.length = 0;
  158. return output;
  159. }
  160. var curr = this.root;
  161. var stack = this._stack;
  162. var idx = 0;
  163. var nearestNList = this._nearstNList;
  164. for (var i = 0; i < N; i++) {
  165. // Allocate
  166. if (!nearestNList[i]) {
  167. nearestNList[i] = {};
  168. }
  169. nearestNList[i].dist = 0;
  170. nearestNList[i].node = null;
  171. }
  172. var currDist = squaredDistance(curr.data, target);
  173. var found = 0;
  174. if (curr.data !== target) {
  175. found++;
  176. this._addNearest(found, currDist, curr);
  177. }
  178. if (target.array[curr.axis] < curr.data.array[curr.axis]) {
  179. // Left first
  180. curr.right && (stack[idx++] = curr.right);
  181. curr.left && (stack[idx++] = curr.left);
  182. }
  183. else {
  184. // Right first
  185. curr.left && (stack[idx++] = curr.left);
  186. curr.right && (stack[idx++] = curr.right);
  187. }
  188. while (idx--) {
  189. curr = stack[idx];
  190. var currDist = target.array[curr.axis] - curr.data.array[curr.axis];
  191. var isLeft = currDist < 0;
  192. var needsCheckOtherSide = false;
  193. currDist = currDist * currDist;
  194. // Intersecting right hyperplane with minDist hypersphere
  195. if (found < N || currDist < nearestNList[found - 1].dist) {
  196. currDist = squaredDistance(curr.data, target);
  197. if (
  198. (found < N || currDist < nearestNList[found - 1].dist)
  199. && curr.data !== target
  200. ) {
  201. if (found < N) {
  202. found++;
  203. }
  204. this._addNearest(found, currDist, curr);
  205. }
  206. needsCheckOtherSide = true;
  207. }
  208. if (isLeft) {
  209. if (needsCheckOtherSide) {
  210. curr.right && (stack[idx++] = curr.right);
  211. }
  212. // Search in the left area
  213. curr.left && (stack[idx++] = curr.left);
  214. }
  215. else {
  216. if (needsCheckOtherSide) {
  217. curr.left && (stack[idx++] = curr.left);
  218. }
  219. // Search the right area
  220. curr.right && (stack[idx++] = curr.right);
  221. }
  222. }
  223. // Copy to output
  224. for (var i = 0; i < found; i++) {
  225. output[i] = nearestNList[i].node.data;
  226. }
  227. output.length = found;
  228. return output;
  229. };
  230. export default KDTree;