Finding MK Average

hard design heap ordered set

Problem

Design a class MKAverage built from two integers m and k. The method addElement(num) appends num to a running stream. The method calculateMKAverage() looks at the last m numbers added: if fewer than m exist it returns -1; otherwise it sorts that window, drops the smallest k and the largest k values, and returns the integer average (floor) of the remaining m − 2k numbers.

Input

m = 3, k = 1 · addElement(3), addElement(1), calculateMKAverage(), addElement(10), calculateMKAverage(), addElement(5), addElement(5), addElement(5), calculateMKAverage()

Output-1, 3, 5

After only 2 adds the window has < 3 numbers, so the query returns -1. With window {3, 1, 10}, dropping the smallest (1) and largest (10) leaves {3}, average 3. Later the last 3 are {5, 5, 5}; trimming one from each end leaves {5}, average 5.

operations (e.g. add(3), avg())

from collections import deque

class MKAverage:
    def __init__(self, m, k):
        self.m, self.k = m, k
        self.q = deque()

    def addElement(self, num):
        self.q.append(num)
        if len(self.q) > self.m:
            self.q.popleft()

    def calculateMKAverage(self):
        if len(self.q) < self.m:
            return -1
        s = sorted(self.q)
        mid = s[self.k:self.m - self.k]
        return sum(mid) // len(mid)

class MKAverage {
  constructor(m, k) {
    this.m = m;
    this.k = k;
    this.q = [];
  }

  addElement(num) {
    this.q.push(num);
    if (this.q.length > this.m) this.q.shift();
  }

  calculateMKAverage() {
    if (this.q.length < this.m) return -1;
    const s = [...this.q].sort((a, b) => a - b);
    const mid = s.slice(this.k, this.m - this.k);
    const sum = mid.reduce((a, b) => a + b, 0);
    return Math.floor(sum / mid.length);
  }
}

class MKAverage {
    private int m, k;
    private Deque<Integer> q = new ArrayDeque<>();

    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
    }

    public void addElement(int num) {
        q.addLast(num);
        if (q.size() > m) q.pollFirst();
    }

    public int calculateMKAverage() {
        if (q.size() < m) return -1;
        List<Integer> s = new ArrayList<>(q);
        Collections.sort(s);
        long sum = 0;
        for (int i = k; i < m - k; i++) sum += s.get(i);
        return (int)(sum / (m - 2 * k));
    }
}

class MKAverage {
    int m, k;
    deque<int> q;
public:
    MKAverage(int m, int k) : m(m), k(k) {}

    void addElement(int num) {
        q.push_back(num);
        if ((int)q.size() > m) q.pop_front();
    }

    int calculateMKAverage() {
        if ((int)q.size() < m) return -1;
        vector<int> s(q.begin(), q.end());
        sort(s.begin(), s.end());
        long sum = 0;
        for (int i = k; i < m - k; i++) sum += s[i];
        return (int)(sum / (m - 2 * k));
    }
};

Explanation

Two things are happening at once. First, only the last m numbers matter — this is a fixed-size sliding window over the stream. A queue handles that perfectly: addElement pushes to the back, and whenever the queue grows past m we pop the oldest from the front. Everything else in the stream is irrelevant to future queries.

Second, a query needs a trimmed mean: sort the window, throw away the smallest k and the largest k, and average the middle m − 2k values (using floor / integer division). If fewer than m numbers have arrived yet, the window is incomplete and the answer is -1.

The version shown above re-sorts the window on every query, which is the clearest way to see the algorithm: it is O(m log m) per calculateMKAverage. That is fine for learning and for moderate input sizes.

To make every query O(log m), the production trick is to keep the window split across three ordered multisets — low (the smallest k), mid (the middle m − 2k), and high (the largest k) — together with a running sum of mid. Each add/evict moves at most a constant number of elements across the boundaries to restore the sizes k / (m−2k) / k, and the answer is just midSum / (m − 2k). The heap/ordered-set tags point at exactly this structure.

Trace the example with m = 3, k = 1. After add(3), add(1) the window is {3, 1} — only two numbers, so the first query returns -1. After add(10) the window is {3, 1, 10}; sorted it is [1, 3, 10], dropping the ends leaves [3], average 3. Three more add(5) calls slide the window to {5, 5, 5} (older values fall off the front); trimming leaves [5], average 5.

Time: O(log m) per op with three ordered multisets (O(m log m) for the simple re-sort shown) Space: O(m)