Finding MK Average LeetCode solution in JavaScript

1
class ArraySegTree {
2
  // Array to perfrom operations on, range query operation, PointUpdate operation
3
  constructor(A, op = (a, b) => a + b, upOp = (a, b) => a + b, opSentinel = 0) {
4
    (this.n = A.length),
5
      (this.t = [...Array(4 * this.n + 1)]),
6
      (this.op = op),
7
      (this.upOp = upOp),
8
      (this.opSentinel = opSentinel);
9
    //root's idx =1
10
    this.build(A, 1, 0, this.n - 1);
11
  }
12
  left = (x) => this.t[2 * x];
13
  right = (x) => this.t[2 * x + 1];
14
  build(A, idx, left, right) {
15
    if (left == right) return (this.t[idx] = A[left]);
16
    let mid = (left + right) >> 1;
17
    this.build(A, 2 * idx, left, mid); //go left
18
    this.build(A, 2 * idx + 1, mid + 1, right); //go right
19
    this.t[idx] = this.op(this.left(idx), this.right(idx)); //merge
20
  }
21
  //just specify l,r on actual queries
22
  //Here queries use the actul indices of the starting array A, so rangeQuery(0,n-1) returns the whole array
23
  rangeQuery = (l, r, tl = 0, tr = this.n - 1, idx = 1) => {
24
    if (l > r) return this.opSentinel;
25
    if (l === tl && r === tr) return this.t[idx];
26
    let mid = (tl + tr) >> 1;
27
    return this.op(
28
      this.rangeQuery(l, Math.min(r, mid), tl, mid, idx * 2),
29
      this.rangeQuery(Math.max(l, mid + 1), r, mid + 1, tr, idx * 2 + 1)
30
    );
31
  };
32
  //just specify arrIdx,newVal on actual pointUpdates
33
  pointUpdate = (arrIdx, newVal, tl = 0, tr = this.n - 1, idx = 1) => {
34
    if (tl == tr) return (this.t[idx] = this.upOp(this.t[idx], newVal));
35
    let mid = (tl + tr) >> 1;
36
    if (arrIdx <= mid) this.pointUpdate(arrIdx, newVal, tl, mid, 2 * idx);
37
    else this.pointUpdate(arrIdx, newVal, mid + 1, tr, 2 * idx + 1);
38
    this.t[idx] = this.op(this.left(idx), this.right(idx));
39
  };
40
  seachPrefixIndex = (x, start = 0, end = this.n - 1) => {
41
    let s = this.rangeQuery(start, end);
42
    if (s < x) return -1;
43
    if (start === end) return start;
44
    let mid = (start + end) >> 1,
45
      left = this.rangeQuery(start, mid);
46
    if (left >= x) return this.seachPrefixIndex(x, start, mid);
47
    else return this.seachPrefixIndex(x - left, mid + 1, end);
48
  };
49
}
50

51
var MKAverage = function (m, k) {
52
  (this.A = []), (this.m = m), (this.k = k), (this.div = m - 2 * k);
53
  this.S1 = new ArraySegTree([...Array(1e5 + 2)].map((d) => 0)); //indices
54
  this.S2 = new ArraySegTree([...Array(1e5 + 2)].map((d) => 0)); //sums
55
};
56

57
/**
58
 * @param {number} num
59
 * @return {void}
60
 */
61
MKAverage.prototype.addElement = function (num) {
62
  this.A.push(num);
63
  this.S1.pointUpdate(num, 1);
64
  this.S2.pointUpdate(num, num);
65
  if (this.A.length > this.m) {
66
    let z = this.A.shift();
67
    this.S1.pointUpdate(z, -1);
68
    this.S2.pointUpdate(z, -z);
69
  }
70
};
71

72
/**
73
 * @return {number}
74
 */
75
MKAverage.prototype.calculateMKAverage = function () {
76
  if (this.A.length < this.m) return -1;
77

78
  let S = this.S1.seachPrefixIndex(this.k + 1),
79
    count1 = this.S1.rangeQuery(0, S),
80
    take = (count1 - this.k) * S;
81

82
  let E = this.S1.seachPrefixIndex(this.m - this.k),
83
    count2 = this.S1.rangeQuery(0, E),
84
    remove = (count2 - this.m + this.k) * E;
85

86
  let midSum = this.S2.rangeQuery(S + 1, E);
87
  return ((take + midSum - remove) / this.div) >> 0;
88
};