kdbush.js 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336
  1. (function (global, factory) {
  2. typeof exports === 'object' && typeof module !== 'undefined' ? module.exports = factory() :
  3. typeof define === 'function' && define.amd ? define(factory) :
  4. (global = typeof globalThis !== 'undefined' ? globalThis : global || self, global.KDBush = factory());
  5. })(this, (function () { 'use strict';
  6. const ARRAY_TYPES = [
  7. Int8Array, Uint8Array, Uint8ClampedArray, Int16Array, Uint16Array,
  8. Int32Array, Uint32Array, Float32Array, Float64Array
  9. ];
  10. /** @typedef {Int8ArrayConstructor | Uint8ArrayConstructor | Uint8ClampedArrayConstructor | Int16ArrayConstructor | Uint16ArrayConstructor | Int32ArrayConstructor | Uint32ArrayConstructor | Float32ArrayConstructor | Float64ArrayConstructor} TypedArrayConstructor */
  11. const VERSION = 1; // serialized format version
  12. const HEADER_SIZE = 8;
  13. class KDBush {
  14. /**
  15. * Creates an index from raw `ArrayBuffer` data.
  16. * @param {ArrayBuffer} data
  17. */
  18. static from(data) {
  19. if (!(data instanceof ArrayBuffer)) {
  20. throw new Error('Data must be an instance of ArrayBuffer.');
  21. }
  22. const [magic, versionAndType] = new Uint8Array(data, 0, 2);
  23. if (magic !== 0xdb) {
  24. throw new Error('Data does not appear to be in a KDBush format.');
  25. }
  26. const version = versionAndType >> 4;
  27. if (version !== VERSION) {
  28. throw new Error(`Got v${version} data when expected v${VERSION}.`);
  29. }
  30. const ArrayType = ARRAY_TYPES[versionAndType & 0x0f];
  31. if (!ArrayType) {
  32. throw new Error('Unrecognized array type.');
  33. }
  34. const [nodeSize] = new Uint16Array(data, 2, 1);
  35. const [numItems] = new Uint32Array(data, 4, 1);
  36. return new KDBush(numItems, nodeSize, ArrayType, data);
  37. }
  38. /**
  39. * Creates an index that will hold a given number of items.
  40. * @param {number} numItems
  41. * @param {number} [nodeSize=64] Size of the KD-tree node (64 by default).
  42. * @param {TypedArrayConstructor} [ArrayType=Float64Array] The array type used for coordinates storage (`Float64Array` by default).
  43. * @param {ArrayBuffer} [data] (For internal use only)
  44. */
  45. constructor(numItems, nodeSize = 64, ArrayType = Float64Array, data) {
  46. if (isNaN(numItems) || numItems < 0) throw new Error(`Unpexpected numItems value: ${numItems}.`);
  47. this.numItems = +numItems;
  48. this.nodeSize = Math.min(Math.max(+nodeSize, 2), 65535);
  49. this.ArrayType = ArrayType;
  50. this.IndexArrayType = numItems < 65536 ? Uint16Array : Uint32Array;
  51. const arrayTypeIndex = ARRAY_TYPES.indexOf(this.ArrayType);
  52. const coordsByteSize = numItems * 2 * this.ArrayType.BYTES_PER_ELEMENT;
  53. const idsByteSize = numItems * this.IndexArrayType.BYTES_PER_ELEMENT;
  54. const padCoords = (8 - idsByteSize % 8) % 8;
  55. if (arrayTypeIndex < 0) {
  56. throw new Error(`Unexpected typed array class: ${ArrayType}.`);
  57. }
  58. if (data && (data instanceof ArrayBuffer)) { // reconstruct an index from a buffer
  59. this.data = data;
  60. this.ids = new this.IndexArrayType(this.data, HEADER_SIZE, numItems);
  61. this.coords = new this.ArrayType(this.data, HEADER_SIZE + idsByteSize + padCoords, numItems * 2);
  62. this._pos = numItems * 2;
  63. this._finished = true;
  64. } else { // initialize a new index
  65. this.data = new ArrayBuffer(HEADER_SIZE + coordsByteSize + idsByteSize + padCoords);
  66. this.ids = new this.IndexArrayType(this.data, HEADER_SIZE, numItems);
  67. this.coords = new this.ArrayType(this.data, HEADER_SIZE + idsByteSize + padCoords, numItems * 2);
  68. this._pos = 0;
  69. this._finished = false;
  70. // set header
  71. new Uint8Array(this.data, 0, 2).set([0xdb, (VERSION << 4) + arrayTypeIndex]);
  72. new Uint16Array(this.data, 2, 1)[0] = nodeSize;
  73. new Uint32Array(this.data, 4, 1)[0] = numItems;
  74. }
  75. }
  76. /**
  77. * Add a point to the index.
  78. * @param {number} x
  79. * @param {number} y
  80. * @returns {number} An incremental index associated with the added item (starting from `0`).
  81. */
  82. add(x, y) {
  83. const index = this._pos >> 1;
  84. this.ids[index] = index;
  85. this.coords[this._pos++] = x;
  86. this.coords[this._pos++] = y;
  87. return index;
  88. }
  89. /**
  90. * Perform indexing of the added points.
  91. */
  92. finish() {
  93. const numAdded = this._pos >> 1;
  94. if (numAdded !== this.numItems) {
  95. throw new Error(`Added ${numAdded} items when expected ${this.numItems}.`);
  96. }
  97. // kd-sort both arrays for efficient search
  98. sort(this.ids, this.coords, this.nodeSize, 0, this.numItems - 1, 0);
  99. this._finished = true;
  100. return this;
  101. }
  102. /**
  103. * Search the index for items within a given bounding box.
  104. * @param {number} minX
  105. * @param {number} minY
  106. * @param {number} maxX
  107. * @param {number} maxY
  108. * @returns {number[]} An array of indices correponding to the found items.
  109. */
  110. range(minX, minY, maxX, maxY) {
  111. if (!this._finished) throw new Error('Data not yet indexed - call index.finish().');
  112. const {ids, coords, nodeSize} = this;
  113. const stack = [0, ids.length - 1, 0];
  114. const result = [];
  115. // recursively search for items in range in the kd-sorted arrays
  116. while (stack.length) {
  117. const axis = stack.pop() || 0;
  118. const right = stack.pop() || 0;
  119. const left = stack.pop() || 0;
  120. // if we reached "tree node", search linearly
  121. if (right - left <= nodeSize) {
  122. for (let i = left; i <= right; i++) {
  123. const x = coords[2 * i];
  124. const y = coords[2 * i + 1];
  125. if (x >= minX && x <= maxX && y >= minY && y <= maxY) result.push(ids[i]);
  126. }
  127. continue;
  128. }
  129. // otherwise find the middle index
  130. const m = (left + right) >> 1;
  131. // include the middle item if it's in range
  132. const x = coords[2 * m];
  133. const y = coords[2 * m + 1];
  134. if (x >= minX && x <= maxX && y >= minY && y <= maxY) result.push(ids[m]);
  135. // queue search in halves that intersect the query
  136. if (axis === 0 ? minX <= x : minY <= y) {
  137. stack.push(left);
  138. stack.push(m - 1);
  139. stack.push(1 - axis);
  140. }
  141. if (axis === 0 ? maxX >= x : maxY >= y) {
  142. stack.push(m + 1);
  143. stack.push(right);
  144. stack.push(1 - axis);
  145. }
  146. }
  147. return result;
  148. }
  149. /**
  150. * Search the index for items within a given radius.
  151. * @param {number} qx
  152. * @param {number} qy
  153. * @param {number} r Query radius.
  154. * @returns {number[]} An array of indices correponding to the found items.
  155. */
  156. within(qx, qy, r) {
  157. if (!this._finished) throw new Error('Data not yet indexed - call index.finish().');
  158. const {ids, coords, nodeSize} = this;
  159. const stack = [0, ids.length - 1, 0];
  160. const result = [];
  161. const r2 = r * r;
  162. // recursively search for items within radius in the kd-sorted arrays
  163. while (stack.length) {
  164. const axis = stack.pop() || 0;
  165. const right = stack.pop() || 0;
  166. const left = stack.pop() || 0;
  167. // if we reached "tree node", search linearly
  168. if (right - left <= nodeSize) {
  169. for (let i = left; i <= right; i++) {
  170. if (sqDist(coords[2 * i], coords[2 * i + 1], qx, qy) <= r2) result.push(ids[i]);
  171. }
  172. continue;
  173. }
  174. // otherwise find the middle index
  175. const m = (left + right) >> 1;
  176. // include the middle item if it's in range
  177. const x = coords[2 * m];
  178. const y = coords[2 * m + 1];
  179. if (sqDist(x, y, qx, qy) <= r2) result.push(ids[m]);
  180. // queue search in halves that intersect the query
  181. if (axis === 0 ? qx - r <= x : qy - r <= y) {
  182. stack.push(left);
  183. stack.push(m - 1);
  184. stack.push(1 - axis);
  185. }
  186. if (axis === 0 ? qx + r >= x : qy + r >= y) {
  187. stack.push(m + 1);
  188. stack.push(right);
  189. stack.push(1 - axis);
  190. }
  191. }
  192. return result;
  193. }
  194. }
  195. /**
  196. * @param {Uint16Array | Uint32Array} ids
  197. * @param {InstanceType<TypedArrayConstructor>} coords
  198. * @param {number} nodeSize
  199. * @param {number} left
  200. * @param {number} right
  201. * @param {number} axis
  202. */
  203. function sort(ids, coords, nodeSize, left, right, axis) {
  204. if (right - left <= nodeSize) return;
  205. const m = (left + right) >> 1; // middle index
  206. // sort ids and coords around the middle index so that the halves lie
  207. // either left/right or top/bottom correspondingly (taking turns)
  208. select(ids, coords, m, left, right, axis);
  209. // recursively kd-sort first half and second half on the opposite axis
  210. sort(ids, coords, nodeSize, left, m - 1, 1 - axis);
  211. sort(ids, coords, nodeSize, m + 1, right, 1 - axis);
  212. }
  213. /**
  214. * Custom Floyd-Rivest selection algorithm: sort ids and coords so that
  215. * [left..k-1] items are smaller than k-th item (on either x or y axis)
  216. * @param {Uint16Array | Uint32Array} ids
  217. * @param {InstanceType<TypedArrayConstructor>} coords
  218. * @param {number} k
  219. * @param {number} left
  220. * @param {number} right
  221. * @param {number} axis
  222. */
  223. function select(ids, coords, k, left, right, axis) {
  224. while (right > left) {
  225. if (right - left > 600) {
  226. const n = right - left + 1;
  227. const m = k - left + 1;
  228. const z = Math.log(n);
  229. const s = 0.5 * Math.exp(2 * z / 3);
  230. const sd = 0.5 * Math.sqrt(z * s * (n - s) / n) * (m - n / 2 < 0 ? -1 : 1);
  231. const newLeft = Math.max(left, Math.floor(k - m * s / n + sd));
  232. const newRight = Math.min(right, Math.floor(k + (n - m) * s / n + sd));
  233. select(ids, coords, k, newLeft, newRight, axis);
  234. }
  235. const t = coords[2 * k + axis];
  236. let i = left;
  237. let j = right;
  238. swapItem(ids, coords, left, k);
  239. if (coords[2 * right + axis] > t) swapItem(ids, coords, left, right);
  240. while (i < j) {
  241. swapItem(ids, coords, i, j);
  242. i++;
  243. j--;
  244. while (coords[2 * i + axis] < t) i++;
  245. while (coords[2 * j + axis] > t) j--;
  246. }
  247. if (coords[2 * left + axis] === t) swapItem(ids, coords, left, j);
  248. else {
  249. j++;
  250. swapItem(ids, coords, j, right);
  251. }
  252. if (j <= k) left = j + 1;
  253. if (k <= j) right = j - 1;
  254. }
  255. }
  256. /**
  257. * @param {Uint16Array | Uint32Array} ids
  258. * @param {InstanceType<TypedArrayConstructor>} coords
  259. * @param {number} i
  260. * @param {number} j
  261. */
  262. function swapItem(ids, coords, i, j) {
  263. swap(ids, i, j);
  264. swap(coords, 2 * i, 2 * j);
  265. swap(coords, 2 * i + 1, 2 * j + 1);
  266. }
  267. /**
  268. * @param {InstanceType<TypedArrayConstructor>} arr
  269. * @param {number} i
  270. * @param {number} j
  271. */
  272. function swap(arr, i, j) {
  273. const tmp = arr[i];
  274. arr[i] = arr[j];
  275. arr[j] = tmp;
  276. }
  277. /**
  278. * @param {number} ax
  279. * @param {number} ay
  280. * @param {number} bx
  281. * @param {number} by
  282. */
  283. function sqDist(ax, ay, bx, by) {
  284. const dx = ax - bx;
  285. const dy = ay - by;
  286. return dx * dx + dy * dy;
  287. }
  288. return KDBush;
  289. }));