Finding MK Average

Hard

Asked by:

10 views

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:

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.
Remove the smallest k elements and the largest k elements from the container.
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:

Input
["MKAverage", "addElement", "addElement", "calculateMKAverage", "addElement", "calculateMKAverage", "addElement", "addElement", "addElement", "calculateMKAverage"]
[[3, 1], [3], [1], [], [10], [], [5], [5], [5], []]
Output
[null, null, null, -1, null, 3, null, null, null, 5]

Explanation
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 <= 10⁵
1 < k*2 < m
1 <= num <= 10⁵
At most 10⁵ calls will be made to addElement and calculateMKAverage.

Try coding it in LeetCode

Solution

Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

When removing the `k` smallest and `k` largest elements, how should duplicates be handled? For example, if the `m` elements are `[1, 2, 2, 3, 4]` and `k=2`, are the two smallest elements considered to be `1` and one of the `2`s?
To confirm, does 'last `m` elements' imply a sliding window? Meaning, once `m` elements are in the stream, each subsequent `addElement` call effectively discards the oldest element from consideration for the MKAverage calculation?
Given that `m` and `num` can both be as large as 10^5, the sum of the middle elements could be up to 10^10, which would overflow a standard 32-bit integer. Should I ensure I use a 64-bit integer type for this summation?
The constraint `k*2 < m` guarantees that the number of elements we average (`m - 2k`) is always at least one. Can I assume this constraint will always hold, meaning I don't need to handle cases of division by zero?
What is the expected behavior of `calculateMKAverage()` if it is called on a newly initialized object, before any elements have been added to the stream?

Brute Force Solution

Approach

The straightforward approach is to maintain a running history of all the numbers added. Each time we need to find the MK average, we'll perform a fresh calculation by looking at only the most recent group of numbers, ordering them, removing the smallest and largest ones, and then averaging what's left.

Here's how the algorithm would work step-by-step:

First, check if you have collected at least 'm' numbers in total. If not, the average cannot be calculated.
If you have enough numbers, take a temporary copy of only the most recent 'm' numbers.
Arrange this temporary collection of 'm' numbers in order from the smallest to the largest.
From this ordered collection, discard the 'k' smallest numbers from the beginning.
Next, discard the 'k' largest numbers from the end of the ordered collection.
You will now be left with only the numbers that were in the middle.
Add up all of these remaining middle numbers to find their total sum.
Finally, divide that total sum by how many numbers are in the middle group to get the MK average.

Code Implementation

class MKAverage:
    def __init__(self, m: int, k: int):
        self.m_value = m
        self.k_value = k
        # A list is used to maintain the full history of numbers in their original order.

        self.number_history = []

    def addElement(self, num: int) -> None:
        self.number_history.append(num)

    def calculateMKAverage(self) -> int:
        # The MK Average is undefined until we have at least 'm' numbers to consider.

        if len(self.number_history) < self.m_value:
            return -1

        recent_m_numbers = self.number_history[-self.m_value:]

        # Sorting is the core of this simple approach to easily find and remove the extremes.

        sorted_recent_numbers = sorted(recent_m_numbers)

        # We slice the sorted list to remove the k smallest and k largest elements.

        numbers_for_average = sorted_recent_numbers[self.k_value:-self.k_value]

        sum_of_middle_elements = sum(numbers_for_average)
        count_of_middle_elements = self.m_value - (2 * self.k_value)

        return sum_of_middle_elements // count_of_middle_elements

Big(O) Analysis

Time Complexity

O(m log m) – The primary computational cost occurs within the calculation of the MK average. For each calculation, a copy of the most recent `m` numbers is made, which takes O(m) time. The most time-consuming step is sorting this temporary collection of `m` numbers, which has a time complexity of O(m log m) using a standard sorting algorithm. All other operations, like slicing and summing the middle elements, are linear in `m` and thus are dominated by the sorting cost. Therefore, the overall time complexity for each calculation simplifies to O(m log m).

Space Complexity

O(N) – The primary driver of space usage is the need to 'maintain a running history of all the numbers added' in a list-like data structure. If N represents the total number of elements added over the lifetime of the object, this history list will grow to a size of N. While the calculation step creates a 'temporary copy' of the most recent 'm' numbers, which uses O(m) space, the persistent history list is the dominant factor. Thus, the overall auxiliary space complexity is O(N).

Optimal Solution

Approach

Instead of repeatedly sorting all the recent numbers every time, the optimal approach keeps them organized into three separate groups: the smallest, the largest, and the middle ones. By constantly keeping these groups balanced as new numbers arrive and old ones depart, the numbers needed for the average are always ready for a quick calculation.

Here's how the algorithm would work step-by-step:

Imagine you have three buckets to hold the last 'm' numbers you've seen. The first bucket is for the 'k' smallest numbers, the third is for the 'k' largest, and the middle bucket holds everything else.
When a new number is added, first figure out which bucket it belongs in by comparing it to the numbers at the boundaries of the buckets, and place it inside.
Placing a new number can make a bucket overflow. To fix this, take the most out-of-place number from the overflowing bucket and move it to its neighboring bucket. For example, if the 'smallest' bucket is too full, its largest number gets moved to the 'middle' bucket.
This might cause a domino effect. If moving a number makes the next bucket too full, repeat the process by moving a number from that bucket to the next, until all three buckets have the correct number of items again.
When the stream of numbers gets long enough, you'll need to remove the oldest number to keep the window at size 'm'. Find which bucket this oldest number is in and take it out.
Removing a number leaves a bucket with a gap. To fix this, pull an appropriate number from a neighboring bucket to fill the spot. This may also cause a domino effect of shuffling that you resolve just like before.
The 'middle' bucket now conveniently holds exactly the numbers you need for the MK average. To find the answer, you just need the sum of the numbers in this middle bucket, which is much faster than sorting the whole collection each time.

Code Implementation

class MKAverage:

    def __init__(self, m: int, k: int):
        self.window_size_m = m
        self.trim_size_k = k
        self.recent_numbers_queue = collections.deque()
        self.smalls_box = []
        self.middles_box = []
        self.bigs_box = []
        self.middle_elements_sum = 0

    def addElement(self, new_element: int) -> None:
        # First, handle removing the oldest element to maintain the m-sized sliding window.
        if len(self.recent_numbers_queue) == self.window_size_m:
            element_to_remove = self.recent_numbers_queue.popleft()
            if self.smalls_box and element_to_remove <= self.smalls_box[-1]:
                self.smalls_box.pop(bisect.bisect_left(self.smalls_box, element_to_remove))
            elif self.bigs_box and element_to_remove >= self.bigs_box[0]:
                self.bigs_box.pop(bisect.bisect_left(self.bigs_box, element_to_remove))
            else:
                self.middles_box.pop(bisect.bisect_left(self.middles_box, element_to_remove))
                self.middle_elements_sum -= element_to_remove
        
        self.recent_numbers_queue.append(new_element)
        
        # The new number is added to a box based on the boundaries of the sorted groups.
        if self.smalls_box and new_element <= self.smalls_box[-1]:
            bisect.insort_left(self.smalls_box, new_element)
        elif self.bigs_box and new_element >= self.bigs_box[0]:
            bisect.insort_left(self.bigs_box, new_element)
        else:
            bisect.insort_left(self.middles_box, new_element)
            self.middle_elements_sum += new_element

        # After addition and removal, the boxes must be rebalanced to their correct sizes.
        while len(self.smalls_box) > self.trim_size_k:
            element_to_move = self.smalls_box.pop()
            bisect.insort_left(self.middles_box, element_to_move)
            self.middle_elements_sum += element_to_move
        while len(self.bigs_box) > self.trim_size_k:
            element_to_move = self.bigs_box.pop(0)
            bisect.insort_left(self.middles_box, element_to_move)
            self.middle_elements_sum += element_to_move

        # Conversely, if boundary boxes are undersized, they are replenished from the middles box.
        while len(self.smalls_box) < self.trim_size_k and self.middles_box:
            element_to_move = self.middles_box.pop(0)
            bisect.insort_left(self.smalls_box, element_to_move)
            self.middle_elements_sum -= element_to_move
        while len(self.bigs_box) < self.trim_size_k and self.middles_box:
            element_to_move = self.middles_box.pop()
            bisect.insort_left(self.bigs_box, element_to_move)
            self.middle_elements_sum -= element_to_move

    def calculateMKAverage(self) -> int:
        if len(self.recent_numbers_queue) < self.window_size_m:
            return -1
        
        middle_box_expected_size = self.window_size_m - 2 * self.trim_size_k
        return self.middle_elements_sum // middle_box_expected_size

Big(O) Analysis

Time Complexity

O(n * log m) – The overall runtime is driven by processing each of the `n` numbers added to the stream. For each new number, the algorithm performs a constant number of operations, such as additions and deletions, on three sorted data structures that together maintain the `m` element window. Using an efficient implementation for these sorted groups, like balanced binary search trees, each operation costs O(log m) time. Since this logarithmic work is performed for every one of the `n` elements, the total time complexity is O(n * log m).

Space Complexity

O(m) – The algorithm's memory usage is determined by the need to store the last 'm' numbers. This is achieved using three data structures, the "buckets," which collectively hold 'm' elements, and an additional queue, also of size 'm', to track the arrival order for removals. Because the size of these data structures is directly proportional to the window size 'm', regardless of how many numbers are processed in total, the auxiliary space complexity is O(m).

Edge Cases

Case	How to Handle
Calculating the average before m elements have been added.	The solution must check the stream size and return -1 as specified in the problem description.
The sliding window contains many or all duplicate values.	Data structures must correctly handle multisets to ensure the k smallest and k largest elements are identified properly, including their duplicates.
The sum of the middle m - 2k elements overflows a standard 32-bit integer.	The running sum must be maintained using a 64-bit integer type to prevent overflow, as the sum can reach up to 10^10.
Maximum constraints are hit for m, k, and the number of calls.	An efficient solution using balanced data structures is required for addElement to achieve logarithmic time complexity and avoid a timeout.
An element is added when the window is full, causing the oldest element to be removed.	A queue must be used to track the window's elements and correctly remove the outgoing element from the sorted data structures.
The input stream is a strictly increasing or decreasing sequence.	The rebalancing logic must efficiently handle the constant migration of elements between the small, middle, and large groups as the window slides.
A newly added or removed element is a boundary value between partitions.	The logic must carefully rebalance elements between the three groups to maintain their size and ordering invariants.
The parameters m and k are at their minimum valid values, such as m=3, k=1.	The partitioning logic must be robust enough to handle the middle group containing just a single element.

Calculating the average before m elements have been added.

How to Handle:

The solution must check the stream size and return -1 as specified in the problem description.

The sliding window contains many or all duplicate values.

How to Handle:

Data structures must correctly handle multisets to ensure the k smallest and k largest elements are identified properly, including their duplicates.

The sum of the middle m - 2k elements overflows a standard 32-bit integer.

How to Handle:

The running sum must be maintained using a 64-bit integer type to prevent overflow, as the sum can reach up to 10^10.

Maximum constraints are hit for m, k, and the number of calls.

How to Handle:

An efficient solution using balanced data structures is required for addElement to achieve logarithmic time complexity and avoid a timeout.

An element is added when the window is full, causing the oldest element to be removed.

How to Handle:

A queue must be used to track the window's elements and correctly remove the outgoing element from the sorted data structures.

The input stream is a strictly increasing or decreasing sequence.

How to Handle:

The rebalancing logic must efficiently handle the constant migration of elements between the small, middle, and large groups as the window slides.

A newly added or removed element is a boundary value between partitions.

How to Handle:

The logic must carefully rebalance elements between the three groups to maintain their size and ordering invariants.

The parameters m and k are at their minimum valid values, such as m=3, k=1.

How to Handle:

The partitioning logic must be robust enough to handle the middle group containing just a single element.