Scale Engineer

Finding MK Average

Hard
#1668Time: `addElement`: O(1) Adding and removing from a queue are constant time operations. `calculateMKAverage`: O(m log m) The dominant operation is sorting the list of `m` elements.Space: O(m) We need a queue to store the last `m` elements. The temporary list in `calculateMKAverage` also takes `O(m)` space.
Patterns
DesignData Stream
Data structures
Heap (Priority Queue)Ordered SetQueue

Prompt

You are given two integers, m and k, and a stream of integers. You are tasked to implement a data structure that calculates the MKAverage for the stream.

The MKAverage can be calculated using these steps:

  1. If the number of the elements in the stream is less than m you should consider the MKAverage to be -1. Otherwise, copy the last m elements of the stream to a separate container.
  2. Remove the smallest k elements and the largest k elements from the container.
  3. Calculate the average value for the rest of the elements rounded down to the nearest integer.

Implement the MKAverage class:

  • MKAverage(int m, int k) Initializes the MKAverage object with an empty stream and the two integers m and k.
  • void addElement(int num) Inserts a new element num into the stream.
  • int calculateMKAverage() Calculates and returns the MKAverage for the current stream rounded down to the nearest integer.

 

Example 1:

MKAverage obj = new MKAverage(3, 1); 
obj.addElement(3);        // current elements are [3]
obj.addElement(1);        // current elements are [3,1]
obj.calculateMKAverage(); // return -1, because m = 3 and only 2 elements exist.
obj.addElement(10);       // current elements are [3,1,10]
obj.calculateMKAverage(); // The last 3 elements are [3,1,10].
                          // After removing smallest and largest 1 element the container will be [3].
                          // The average of [3] equals 3/1 = 3, return 3
obj.addElement(5);        // current elements are [3,1,10,5]
obj.addElement(5);        // current elements are [3,1,10,5,5]
obj.addElement(5);        // current elements are [3,1,10,5,5,5]
obj.calculateMKAverage(); // The last 3 elements are [5,5,5].
                          // After removing smallest and largest 1 element the container will be [5].
                          // The average of [5] equals 5/1 = 5, return 5

 

Constraints:

  • 3 <= m <= 105
  • 1 < k*2 < m
  • 1 <= num <= 105
  • At most 105 calls will be made to addElement and calculateMKAverage.

Approaches

3 approaches with complexity analysis and trade-offs.

This approach directly simulates the process described in the problem. It maintains a sliding window of the last m elements using a queue. For each calculateMKAverage call, it copies these m elements into a temporary list, sorts the list, and then calculates the average of the middle elements. While simple to understand, its performance suffers due to the repeated sorting.

Algorithm

  • Initialization: In the constructor, initialize a Queue (e.g., LinkedList) to store the stream elements, and store the values of m and k.
  • addElement(num):
    • Add the new element num to the tail of the queue.
    • If the size of the queue is greater than m, remove the element from the head of the queue (poll()). This maintains a sliding window of the last m elements.
  • calculateMKAverage():
    • Check if the number of elements in the queue is less than m. If it is, return -1.
    • Create a temporary ArrayList and copy all elements from the queue into it.
    • Sort the temporary list using a standard sorting algorithm (e.g., Collections.sort()).
    • Initialize a variable sum to zero.
    • Iterate through the sorted list from index k to m - 1 - k (inclusive) and add each element to sum.
    • Calculate the average by dividing sum by the count of middle elements, which is m - 2*k.
    • Return the result, which is automatically floored due to integer division.

Walkthrough

The core idea is to use a Queue data structure to keep track of the last m elements added to the stream. This naturally handles the sliding window requirement.

When addElement(num) is called, the new number is added to the queue. To ensure the queue only contains the last m elements, we check its size after adding. If it exceeds m, we remove the oldest element, which is at the front of the queue.

When calculateMKAverage() is called, the main work begins. First, a size check ensures we have enough elements. If not, -1 is returned. Otherwise, the current m elements from the queue are copied into a new list. This list is then sorted, which is the most time-consuming step (O(m log m)). After sorting, the smallest k elements are at the beginning of the list, and the largest k are at the end. We can simply ignore them by summing the elements from index k up to (but not including) index m - k. The final average is this sum divided by the number of elements considered (m - 2*k).

import java.util.ArrayList;import java.util.Collections;import java.util.LinkedList;import java.util.List;import java.util.Queue; class MKAverage {    private int m, k;    private Queue<Integer> stream;     public MKAverage(int m, int k) {        this.m = m;        this.k = k;        this.stream = new LinkedList<>();    }        public void addElement(int num) {        stream.offer(num);        if (stream.size() > m) {            stream.poll();        }    }        public int calculateMKAverage() {        if (stream.size() < m) {            return -1;        }         List<Integer> window = new ArrayList<>(stream);        Collections.sort(window);         long sum = 0;        for (int i = k; i < m - k; i++) {            sum += window.get(i);        }         return (int) (sum / (m - 2 * k));    }}

Complexity

Time

`addElement`: O(1) Adding and removing from a queue are constant time operations. `calculateMKAverage`: O(m log m) The dominant operation is sorting the list of `m` elements.

Space

O(m) We need a queue to store the last `m` elements. The temporary list in `calculateMKAverage` also takes `O(m)` space.

Trade-offs

Pros

  • Easy to understand and implement.

  • The logic directly follows the problem statement.

Cons

  • The calculateMKAverage method is inefficient, with a time complexity of O(m log m).

  • For streams with frequent average calculations, this approach can be very slow, potentially leading to a 'Time Limit Exceeded' error in a competitive programming context.

Solutions

class MKAverage { private int m , k ; private long s ; private int size1 , size3 ; private Deque < Integer > q = new ArrayDeque <>(); private TreeMap < Integer , Integer > lo = new TreeMap <>(); private TreeMap < Integer , Integer > mid = new TreeMap <>(); private TreeMap < Integer , Integer > hi = new TreeMap <>(); public MKAverage ( int m , int k ) { this . m = m ; this . k = k ; } public void addElement ( int num ) { if ( lo . isEmpty () || num <= lo . lastKey ()) { lo . merge ( num , 1 , Integer: : sum ); ++ size1 ; } else if ( hi . isEmpty () || num >= hi . firstKey ()) { hi . merge ( num , 1 , Integer: : sum ); ++ size3 ; } else { mid . merge ( num , 1 , Integer: : sum ); s += num ; } q . offer ( num ); if ( q . size () > m ) { int x = q . poll (); if ( lo . containsKey ( x )) { if ( lo . merge ( x , - 1 , Integer: : sum ) == 0 ) { lo . remove ( x ); } -- size1 ; } else if ( hi . containsKey ( x )) { if ( hi . merge ( x , - 1 , Integer: : sum ) == 0 ) { hi . remove ( x ); } -- size3 ; } else { if ( mid . merge ( x , - 1 , Integer: : sum ) == 0 ) { mid . remove ( x ); } s -= x ; } } for (; size1 > k ; -- size1 ) { int x = lo . lastKey (); if ( lo . merge ( x , - 1 , Integer: : sum ) == 0 ) { lo . remove ( x ); } mid . merge ( x , 1 , Integer: : sum ); s += x ; } for (; size3 > k ; -- size3 ) { int x = hi . firstKey (); if ( hi . merge ( x , - 1 , Integer: : sum ) == 0 ) { hi . remove ( x ); } mid . merge ( x , 1 , Integer: : sum ); s += x ; } for (; size1 < k && ! mid . isEmpty (); ++ size1 ) { int x = mid . firstKey (); if ( mid . merge ( x , - 1 , Integer: : sum ) == 0 ) { mid . remove ( x ); } s -= x ; lo . merge ( x , 1 , Integer: : sum ); } for (; size3 < k && ! mid . isEmpty (); ++ size3 ) { int x = mid . lastKey (); if ( mid . merge ( x , - 1 , Integer: : sum ) == 0 ) { mid . remove ( x ); } s -= x ; hi . merge ( x , 1 , Integer: : sum ); } } public int calculateMKAverage () { return q . size () < m ? - 1 : ( int ) ( s / ( q . size () - k * 2 )); } } /** * Your MKAverage object will be instantiated and called as such: * MKAverage obj = new MKAverage(m, k); * obj.addElement(num); * int param_2 = obj.calculateMKAverage(); */

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next →

Next problem