KDTree.js 6.0 KB

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