PrefixIntervalTree.js.flow 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244
  1. /**
  2. * Copyright (c) 2013-present, Facebook, Inc.
  3. *
  4. * This source code is licensed under the MIT license found in the
  5. * LICENSE file in the root directory of this source tree.
  6. *
  7. * @providesModule PrefixIntervalTree
  8. * @flow
  9. * @typechecks
  10. */
  11. 'use strict';
  12. const invariant = require("./invariant");
  13. const parent = node => Math.floor(node / 2);
  14. const Int32Array = global.Int32Array || function (size: number): Array<number> {
  15. const xs = [];
  16. for (let i = size - 1; i >= 0; --i) {
  17. xs[i] = 0;
  18. }
  19. return xs;
  20. };
  21. /**
  22. * Computes the next power of 2 after or equal to x.
  23. */
  24. function ceilLog2(x: number): number {
  25. let y = 1;
  26. while (y < x) {
  27. y *= 2;
  28. }
  29. return y;
  30. }
  31. /**
  32. * A prefix interval tree stores an numeric array and the partial sums of that
  33. * array. It is optimized for updating the values of the array without
  34. * recomputing all of the partial sums.
  35. *
  36. * - O(ln n) update
  37. * - O(1) lookup
  38. * - O(ln n) compute a partial sum
  39. * - O(n) space
  40. *
  41. * Note that the sequence of partial sums is one longer than the array, so that
  42. * the first partial sum is always 0, and the last partial sum is the sum of the
  43. * entire array.
  44. */
  45. class PrefixIntervalTree {
  46. /**
  47. * Number of elements in the array
  48. */
  49. _size: number;
  50. /**
  51. * Half the size of the heap. It is also the number of non-leaf nodes, and the
  52. * index of the first element in the heap. Always a power of 2.
  53. */
  54. _half: number;
  55. /**
  56. * Binary heap
  57. */
  58. _heap: Int32Array;
  59. constructor(xs: Array<number>) {
  60. this._size = xs.length;
  61. this._half = ceilLog2(this._size);
  62. this._heap = new Int32Array(2 * this._half);
  63. let i;
  64. for (i = 0; i < this._size; ++i) {
  65. this._heap[this._half + i] = xs[i];
  66. }
  67. for (i = this._half - 1; i > 0; --i) {
  68. this._heap[i] = this._heap[2 * i] + this._heap[2 * i + 1];
  69. }
  70. }
  71. static uniform(size: number, initialValue: number): PrefixIntervalTree {
  72. const xs = [];
  73. for (let i = size - 1; i >= 0; --i) {
  74. xs[i] = initialValue;
  75. }
  76. return new PrefixIntervalTree(xs);
  77. }
  78. static empty(size: number): PrefixIntervalTree {
  79. return PrefixIntervalTree.uniform(size, 0);
  80. }
  81. set(index: number, value: number): void {
  82. invariant(0 <= index && index < this._size, 'Index out of range %s', index);
  83. let node = this._half + index;
  84. this._heap[node] = value;
  85. node = parent(node);
  86. for (; node !== 0; node = parent(node)) {
  87. this._heap[node] = this._heap[2 * node] + this._heap[2 * node + 1];
  88. }
  89. }
  90. get(index: number): number {
  91. invariant(0 <= index && index < this._size, 'Index out of range %s', index);
  92. const node = this._half + index;
  93. return this._heap[node];
  94. }
  95. getSize(): number {
  96. return this._size;
  97. }
  98. /**
  99. * Returns the sum get(0) + get(1) + ... + get(end - 1).
  100. */
  101. sumUntil(end: number): number {
  102. invariant(0 <= end && end < this._size + 1, 'Index out of range %s', end);
  103. if (end === 0) {
  104. return 0;
  105. }
  106. let node = this._half + end - 1;
  107. let sum = this._heap[node];
  108. for (; node !== 1; node = parent(node)) {
  109. if (node % 2 === 1) {
  110. sum += this._heap[node - 1];
  111. }
  112. }
  113. return sum;
  114. }
  115. /**
  116. * Returns the sum get(0) + get(1) + ... + get(inclusiveEnd).
  117. */
  118. sumTo(inclusiveEnd: number): number {
  119. invariant(0 <= inclusiveEnd && inclusiveEnd < this._size, 'Index out of range %s', inclusiveEnd);
  120. return this.sumUntil(inclusiveEnd + 1);
  121. }
  122. /**
  123. * Returns the sum get(begin) + get(begin + 1) + ... + get(end - 1).
  124. */
  125. sum(begin: number, end: number): number {
  126. invariant(begin <= end, 'Begin must precede end');
  127. return this.sumUntil(end) - this.sumUntil(begin);
  128. }
  129. /**
  130. * Returns the smallest i such that 0 <= i <= size and sumUntil(i) <= t, or
  131. * -1 if no such i exists.
  132. */
  133. greatestLowerBound(t: number): number {
  134. if (t < 0) {
  135. return -1;
  136. }
  137. let node = 1;
  138. if (this._heap[node] <= t) {
  139. return this._size;
  140. }
  141. while (node < this._half) {
  142. const leftSum = this._heap[2 * node];
  143. if (t < leftSum) {
  144. node = 2 * node;
  145. } else {
  146. node = 2 * node + 1;
  147. t -= leftSum;
  148. }
  149. }
  150. return node - this._half;
  151. }
  152. /**
  153. * Returns the smallest i such that 0 <= i <= size and sumUntil(i) < t, or
  154. * -1 if no such i exists.
  155. */
  156. greatestStrictLowerBound(t: number): number {
  157. if (t <= 0) {
  158. return -1;
  159. }
  160. let node = 1;
  161. if (this._heap[node] < t) {
  162. return this._size;
  163. }
  164. while (node < this._half) {
  165. const leftSum = this._heap[2 * node];
  166. if (t <= leftSum) {
  167. node = 2 * node;
  168. } else {
  169. node = 2 * node + 1;
  170. t -= leftSum;
  171. }
  172. }
  173. return node - this._half;
  174. }
  175. /**
  176. * Returns the smallest i such that 0 <= i <= size and t <= sumUntil(i), or
  177. * size + 1 if no such i exists.
  178. */
  179. leastUpperBound(t: number): number {
  180. return this.greatestStrictLowerBound(t) + 1;
  181. }
  182. /**
  183. * Returns the smallest i such that 0 <= i <= size and t < sumUntil(i), or
  184. * size + 1 if no such i exists.
  185. */
  186. leastStrictUpperBound(t: number): number {
  187. return this.greatestLowerBound(t) + 1;
  188. }
  189. }
  190. module.exports = PrefixIntervalTree;