LeetCode Solutions Blog

Finding MK Average

Number: 1953

Difficulty: Hard

Paid? No

Companies: Google

Problem Description

Design a data structure that supports two operations on a stream of integers. It must maintain a sliding window consisting of the last m elements added. When requested, the data structure computes the MKAverage by removing the smallest k and largest k elements from the window and returning the floor of the average of the remaining numbers. If there are fewer than m elements in the stream, return -1.

Key Insights

Use three separate containers to partition the current m elements:
- One container (left) holds the smallest k elements.
- One container (right) holds the largest k elements.
- The middle container (mid) holds the remaining m - 2k elements and maintains a running sum.
Employ a sliding window strategy: when a new element is added and the window is full, remove the oldest element and rebalance the containers.
Balanced data structures (e.g., balanced BSTs or SortedLists) allow insertion/removal in logarithmic time.
Rebalancing the three groups after each insertion/removal is the key to correctly maintaining the smallest and largest k elements.

Space and Time Complexity

Time Complexity: O(log m) per addElement operation due to insertions and removals on balanced trees. Space Complexity: O(m) for storing the sliding window and the elements across the three containers.

Solution

We maintain three ordered sets (or SortedLists) named left, mid, and right. Additionally, we store the m most recent stream elements in a queue for the sliding window. The left set contains the smallest k elements, the right set contains the largest k elements, and the mid set contains the remaining elements whose sum is tracked. When adding a new element:

Insert the new element into one of the three containers based on its value.
Rebalance the containers if the size constraints (left and right must always have k elements) are violated.
If the overall number of elements exceeds m, remove the oldest element from the appropriate container and then rebalance. When calculating MKAverage, if fewer than m elements have been added, return -1; otherwise, return the floor of the average computed from the mid set.

Code Solutions

from collections import deque
from sortedcontainers import SortedList

class MKAverage:
    def __init__(self, m: int, k: int):
        # Store configuration parameters
        self.m = m
        self.k = k
        # Queue to hold the sliding window of last m elements
        self.queue = deque()
        # Three sorted lists for small (left), middle (mid) and large (right) parts.
        self.left = SortedList()  # smallest k elements
        self.mid = SortedList()   # middle m - 2k elements
        self.right = SortedList() # largest k elements
        # Running sum for elements in mid
        self.midSum = 0

    def addElement(self, num: int) -> None:
        # Add new element to the sliding window queue.
        self.queue.append(num)
        
        # If left is not empty and num is less than or equal to the largest in left, add it there.
        if self.left and num <= self.left[-1]:
            self.left.add(num)
        # Else if right is not empty and num is greater than or equal to the smallest in right, add it there.
        elif self.right and num >= self.right[0]:
            self.right.add(num)
        else:
            # Otherwise, it belongs to mid
            self.mid.add(num)
            self.midSum += num

        # Rebalance the three parts so that left and right have exactly k elements each
        self._rebalance_after_add()

        # If we have more than m elements in the stream, remove the oldest element
        if len(self.queue) > self.m:
            remove_val = self.queue.popleft()
            # Remove the outdated element from its respective container
            if remove_val in self.left:
                self.left.remove(remove_val)
            elif remove_val in self.right:
                self.right.remove(remove_val)
            else:
                self.mid.remove(remove_val)
                self.midSum -= remove_val
            # Rebalance after removal
            self._rebalance_after_remove()
    
    def calculateMKAverage(self) -> int:
        # If not enough elements are in the stream, return -1.
        if len(self.queue) < self.m:
            return -1
        # The MKAverage is the floor of the average of the middle collection.
        return self.midSum // len(self.mid)
    
    def _rebalance_after_add(self):
        # Move elements from left -> mid if left size exceeds k.
        while len(self.left) > self.k:
            # Remove the largest element from left and add it to mid.
            val = self.left.pop(-1)
            self.mid.add(val)
            self.midSum += val
        # Move elements from mid -> left if left size is less than k.
        while len(self.left) < self.k and self.mid:
            val = self.mid.pop(0)
            self.midSum -= val
            self.left.add(val)
        
        # Move elements from right -> mid if right size exceeds k.
        while len(self.right) > self.k:
            # Remove the smallest element from right and add it to mid.
            val = self.right.pop(0)
            self.mid.add(val)
            self.midSum += val
        # Move elements from mid -> right if right size is less than k.
        while len(self.right) < self.k and self.mid:
            val = self.mid.pop(-1)
            self.midSum -= val
            self.right.add(val)
    
    def _rebalance_after_remove(self):
        # After removal, re-establish the sizes for left and right from mid.
        while len(self.left) < self.k and self.mid:
            # Move smallest element in mid to left.
            val = self.mid.pop(0)
            self.midSum -= val
            self.left.add(val)
        while len(self.left) > self.k:
            # Move largest element from left back to mid.
            val = self.left.pop(-1)
            self.mid.add(val)
            self.midSum += val
        
        while len(self.right) < self.k and self.mid:
            # Move largest element from mid to right.
            val = self.mid.pop(-1)
            self.midSum -= val
            self.right.add(val)
        while len(self.right) > self.k:
            # Move smallest element from right back to mid.
            val = self.right.pop(0)
            self.mid.add(val)
            self.midSum += val

# Example usage:
# obj = MKAverage(3, 1)
# obj.addElement(3)
# obj.addElement(1)
# print(obj.calculateMKAverage()) # Output: -1 because not enough elements exist.
# obj.addElement(10)
# print(obj.calculateMKAverage()) # Output: 3
# obj.addElement(5)
# obj.addElement(5)
# obj.addElement(5)
# print(obj.calculateMKAverage()) # Output: 5