Finding MK Average

class MKAverage {
    int m, k;
    deque<int> q;
    multiset<int> low, mid, high;
    long long sumMid = 0;
    void add(int x) {
        if ((int)low.size() < k) low.insert(x);
        else if (x <= *low.rbegin()) low.insert(x);
        else if ((int)high.size() < k) high.insert(x);
        else if (x >= *high.begin()) high.insert(x);
        else { mid.insert(x); sumMid += x; }
        rebalance();
    }
    void remove(int x) {
        if (low.count(x)) low.erase(low.find(x));
        else if (high.count(x)) high.erase(high.find(x));
        else { mid.erase(mid.find(x)); sumMid -= x; }
        rebalance();
    }
    void rebalance() {
        while ((int)low.size() > k) {
            int v = *low.rbegin(); low.erase(prev(low.end()));
            mid.insert(v); sumMid += v;
        }
        while ((int)low.size() < k && !mid.empty()) {
            int v = *mid.begin(); mid.erase(mid.begin()); sumMid -= v;
            low.insert(v);
        }
        while ((int)high.size() > k) {
            int v = *high.begin(); high.erase(high.begin());
            mid.insert(v); sumMid += v;
        }
        while ((int)high.size() < k && !mid.empty()) {
            int v = *mid.rbegin(); mid.erase(prev(mid.end())); sumMid -= v;
            high.insert(v);
        }
    }
public:
    MKAverage(int m_, int k_) : m(m_), k(k_) {}
    void addElement(int num) {
        q.push_back(num);
        if ((int)q.size() <= m) {
            add(num);
            return;
        }
        int out = q.front(); q.pop_front();
        remove(out);
        add(num);
    }
    int calculateMKAverage() {
        if ((int)q.size() < m) return -1;
        return (int)(sumMid / (long long)(m - 2 * k));
    }
};

// Fenwick-tree variant: one BIT over counts, one BIT over value-sums.
// addElement maintains a sliding window of last m values; calculateMKAverage
// asks the BIT for the sum of values whose rank lies in (k, m-k].

const MAX_VAL = 100001

type MKAverage struct {
    m, k     int
    q        []int
    cnt, sum []int64
    total    int64
}

func Constructor(m int, k int) MKAverage {
    return MKAverage{
        m:   m,
        k:   k,
        cnt: make([]int64, MAX_VAL+1),
        sum: make([]int64, MAX_VAL+1),
    }
}

func (b *MKAverage) bitUpdate(i int, dc, ds int64) {
    for ; i <= MAX_VAL; i += i & -i {
        b.cnt[i] += dc
        b.sum[i] += ds
    }
}

// prefixSumOfBottomN: sum of the smallest n values currently in the window.
func (b *MKAverage) prefixSumOfBottomN(n int) int64 {
    var pos, cumCnt int64 = 0, 0
    var cumSum int64 = 0
    target := int64(n)
    // walk the Fenwick tree from high to low to find boundary
    step := 1
    for step*2 <= MAX_VAL {
        step *= 2
    }
    for ; step > 0; step >>= 1 {
        next := pos + int64(step)
        if next <= MAX_VAL && cumCnt+b.cnt[next] <= target {
            cumCnt += b.cnt[next]
            cumSum += b.sum[next]
            pos = next
        }
    }
    // pos is the largest index whose prefix count <= target; add residual from pos+1
    remaining := target - cumCnt
    if remaining > 0 {
        cumSum += remaining * (pos + 1)
    }
    return cumSum
}

func (b *MKAverage) AddElement(num int) {
    b.q = append(b.q, num)
    b.bitUpdate(num, 1, int64(num))
    b.total += int64(num)
    if len(b.q) > b.m {
        out := b.q[0]
        b.q = b.q[1:]
        b.bitUpdate(out, -1, -int64(out))
        b.total -= int64(out)
    }
}

func (b *MKAverage) CalculateMKAverage() int {
    if len(b.q) < b.m {
        return -1
    }
    lowSum := b.prefixSumOfBottomN(b.k)
    highCutSum := b.prefixSumOfBottomN(b.m - b.k)
    midSum := highCutSum - lowSum
    return int(midSum / int64(b.m-2*b.k))
}

from collections import deque
from sortedcontainers import SortedList

class MKAverage:
    def __init__(self, m: int, k: int):
        self.m = m
        self.k = k
        self.q: deque = deque()
        self.low = SortedList()
        self.mid = SortedList()
        self.high = SortedList()
        self.sum_mid = 0

    def _add(self, x: int) -> None:
        if len(self.low) < self.k:
            self.low.add(x)
        elif x <= self.low[-1]:
            self.low.add(x)
        elif len(self.high) < self.k:
            self.high.add(x)
        elif x >= self.high[0]:
            self.high.add(x)
        else:
            self.mid.add(x)
            self.sum_mid += x
        self._rebalance()

    def _remove(self, x: int) -> None:
        if x in self.low:
            self.low.remove(x)
        elif x in self.high:
            self.high.remove(x)
        else:
            self.mid.remove(x)
            self.sum_mid -= x
        self._rebalance()

    def _rebalance(self) -> None:
        while len(self.low) > self.k:
            v = self.low.pop()
            self.mid.add(v)
            self.sum_mid += v
        while len(self.low) < self.k and self.mid:
            v = self.mid.pop(0)
            self.sum_mid -= v
            self.low.add(v)
        while len(self.high) > self.k:
            v = self.high.pop(0)
            self.mid.add(v)
            self.sum_mid += v
        while len(self.high) < self.k and self.mid:
            v = self.mid.pop()
            self.sum_mid -= v
            self.high.add(v)

    def addElement(self, num: int) -> None:
        self.q.append(num)
        if len(self.q) <= self.m:
            self._add(num)
            return
        out = self.q.popleft()
        self._remove(out)
        self._add(num)

    def calculateMKAverage(self) -> int:
        if len(self.q) < self.m:
            return -1
        return self.sum_mid // (self.m - 2 * self.k)

// Fenwick-tree variant: one BIT over counts, one BIT over value-sums.
// Sum of values whose rank lies in (k, m-k] is the trimmed sum.
const MAX_VAL = 100001;

class MKAverage {
    constructor(m, k) {
        this.m = m;
        this.k = k;
        this.q = [];
        this.cnt = new Array(MAX_VAL + 1).fill(0);
        this.sum = new Array(MAX_VAL + 1).fill(0);
    }

    bitUpdate(i, dc, ds) {
        for (; i <= MAX_VAL; i += i & -i) {
            this.cnt[i] += dc;
            this.sum[i] += ds;
        }
    }

    // Sum of the smallest n values currently in the window.
    prefixSumOfBottomN(n) {
        let pos = 0, cumCnt = 0, cumSum = 0;
        let step = 1;
        while (step * 2 <= MAX_VAL) step *= 2;
        for (; step > 0; step >>= 1) {
            const next = pos + step;
            if (next <= MAX_VAL && cumCnt + this.cnt[next] <= n) {
                cumCnt += this.cnt[next];
                cumSum += this.sum[next];
                pos = next;
            }
        }
        const remaining = n - cumCnt;
        if (remaining > 0) cumSum += remaining * (pos + 1);
        return cumSum;
    }

    addElement(num) {
        this.q.push(num);
        this.bitUpdate(num, 1, num);
        if (this.q.length > this.m) {
            const out = this.q.shift();
            this.bitUpdate(out, -1, -out);
        }
    }

    calculateMKAverage() {
        if (this.q.length < this.m) return -1;
        const lowSum = this.prefixSumOfBottomN(this.k);
        const highCutSum = this.prefixSumOfBottomN(this.m - this.k);
        const midSum = highCutSum - lowSum;
        return Math.floor(midSum / (this.m - 2 * this.k));
    }
}

class MKAverage {
    private static final int MAX_VAL = 100001;
    private int m, k;
    private Deque<Integer> q;
    private long[] cnt;
    private long[] sum;

    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
        this.q = new ArrayDeque<>();
        this.cnt = new long[MAX_VAL + 1];
        this.sum = new long[MAX_VAL + 1];
    }

    private void bitUpdate(int i, long dc, long ds) {
        for (; i <= MAX_VAL; i += i & -i) {
            cnt[i] += dc;
            sum[i] += ds;
        }
    }

    private long prefixSumOfBottomN(int n) {
        long pos = 0, cumCnt = 0, cumSum = 0;
        int step = 1;
        while (step * 2 <= MAX_VAL) step *= 2;
        for (; step > 0; step >>= 1) {
            long next = pos + step;
            if (next <= MAX_VAL && cumCnt + cnt[(int) next] <= n) {
                cumCnt += cnt[(int) next];
                cumSum += sum[(int) next];
                pos = next;
            }
        }
        long remaining = n - cumCnt;
        if (remaining > 0) cumSum += remaining * (pos + 1);
        return cumSum;
    }

    public void addElement(int num) {
        q.offer(num);
        bitUpdate(num, 1, num);
        if (q.size() > m) {
            int out = q.poll();
            bitUpdate(out, -1, -out);
        }
    }

    public int calculateMKAverage() {
        if (q.size() < m) return -1;
        long lowSum = prefixSumOfBottomN(k);
        long highCutSum = prefixSumOfBottomN(m - k);
        long midSum = highCutSum - lowSum;
        return (int) (midSum / (m - 2L * k));
    }
}

// Fenwick-tree variant: one BIT over counts, one BIT over value-sums.
// Sum of values whose rank lies in (k, m-k] is the trimmed sum.
const MAX_VAL = 100001;

class MKAverage {
    private m: number;
    private k: number;
    private q: number[] = [];
    private cnt: number[];
    private sum: number[];

    constructor(m: number, k: number) {
        this.m = m;
        this.k = k;
        this.cnt = new Array(MAX_VAL + 1).fill(0);
        this.sum = new Array(MAX_VAL + 1).fill(0);
    }

    private bitUpdate(i: number, dc: number, ds: number): void {
        for (; i <= MAX_VAL; i += i & -i) {
            this.cnt[i] += dc;
            this.sum[i] += ds;
        }
    }

    // Sum of the smallest n values currently in the window.
    private prefixSumOfBottomN(n: number): number {
        let pos = 0, cumCnt = 0, cumSum = 0;
        let step = 1;
        while (step * 2 <= MAX_VAL) step *= 2;
        for (; step > 0; step >>= 1) {
            const next = pos + step;
            if (next <= MAX_VAL && cumCnt + this.cnt[next] <= n) {
                cumCnt += this.cnt[next];
                cumSum += this.sum[next];
                pos = next;
            }
        }
        const remaining = n - cumCnt;
        if (remaining > 0) cumSum += remaining * (pos + 1);
        return cumSum;
    }

    addElement(num: number): void {
        this.q.push(num);
        this.bitUpdate(num, 1, num);
        if (this.q.length > this.m) {
            const out = this.q.shift()!;
            this.bitUpdate(out, -1, -out);
        }
    }

    calculateMKAverage(): number {
        if (this.q.length < this.m) return -1;
        const lowSum = this.prefixSumOfBottomN(this.k);
        const highCutSum = this.prefixSumOfBottomN(this.m - this.k);
        const midSum = highCutSum - lowSum;
        return Math.floor(midSum / (this.m - 2 * this.k));
    }
}