Update 1825-finding-mk-average.js

Author: everthis

1825-finding-mk-average.js

@@ -1,3 +1,79 @@
+class BIT {
+  constructor(n) {
+    this.arr = new Array(n + 1).fill(0)
+  }
+
+  update(i, v) {
+    while (i < this.arr.length) {
+      this.arr[i] += v
+      i += i & -i
+    }
+  }
+
+  prefixSum(i) {
+    let res = 0
+    while (i > 0) {
+      res += this.arr[i]
+      i -= i & -i
+    }
+    return res
+  }
+
+  upper_bound(v) {
+    const n = Math.floor(Math.log2(this.arr.length))
+    let pos = 0
+    let s = 0
+    for (let i = n; i >= 0; i--) {
+      if (
+        pos + (1 << i) < this.arr.length &&
+        s + this.arr[pos + (1 << i)] <= v
+      ) {
+        pos += 1 << i
+        s += this.arr[pos]
+      }
+    }
+    return pos + 1
+  }
+}
+
+class MKAverage {
+  constructor(m, k) {
+    this.m = m
+    this.k = k
+    this.cnt = new BIT(10 ** 5)
+    this.bit = new BIT(10 ** 5)
+    this.q = []
+  }
+
+  addElement(num) {
+    this.q.push(num)
+    this.cnt.update(num, 1)
+    this.bit.update(num, num)
+    if (this.q.length > this.m) {
+      const x = this.q.shift()
+      this.cnt.update(x, -1)
+      this.bit.update(x, -x)
+    }
+  }
+
+  calculateMKAverage() {
+    if (this.q.length < this.m) {
+      return -1
+    }
+    const l = this.k
+    const r = this.m - this.k
+    const i = this.cnt.upper_bound(l)
+    const j = this.cnt.upper_bound(r)
+    let ans = this.bit.prefixSum(j) - this.bit.prefixSum(i)
+    ans += (this.cnt.prefixSum(i) - l) * i
+    ans -= (this.cnt.prefixSum(j) - r) * j
+    return Math.floor(ans / (this.m - 2 * this.k))
+  }
+}
+
+// another
+
+
 /**
  * @param {number} m
  * @param {number} k