Hard

Finding MK Average

LAB

Simulate MKAverage over a stream. For each calculate operation, average the last m values after removing the k smallest and k largest, rounded down.

EXAMPLES

Example 1

Input
{
  "m": 3,
  "k": 1,
  "operations": [
    "add",
    "add",
    "calc",
    "add",
    "calc"
  ],
  "values": [
    3,
    1,
    0,
    10,
    0
  ]
}

Output
[
  -1,
  3
]

FUNCTION SHAPE

m: intk: intoperations: stringArrayvalues: intArray→intArray

SOLUTION NOTE

Extension of median finding with 3 SortedLists. Maintain k smallest in left, k largest in right, middle contains the rest. Track middle_sum for O(1) average calculation.

Reveal reference solution +

pythonREFERENCE

from sortedcontainers import SortedList

class MKAverage:
    def __init__(self, m: int, k: int):
        self.m, self.k = m, k
        self.left = SortedList()    # k smallest
        self.middle = SortedList()  # m - 2k middle elements
        self.right = SortedList()   # k largest
        self.middle_sum = 0
        self.stream = []

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

        # Remove element if window exceeds m
        if len(self.stream) > self.m:
            rm = self.stream[-self.m - 1]
            if rm in self.left:
                self.left.remove(rm)
            elif rm in self.middle:
                self.middle.remove(rm)
                self.middle_sum -= rm
            else:
                self.right.remove(rm)

        # Add to left, bubble up through middle to right
        self.left.add(num)
        if len(self.left) > self.k:
            val = self.left.pop(-1)
            self.middle.add(val)
            self.middle_sum += val
        if len(self.middle) > self.m - 2 * self.k:
            val = self.middle.pop(-1)
            self.middle_sum -= val
            self.right.add(val)

        # Rebalance: ensure left has k, middle has m-2k
        while len(self.left) < self.k and self.middle:
            val = self.middle.pop(0)
            self.middle_sum -= val
            self.left.add(val)
        while len(self.middle) < self.m - 2 * self.k and self.right:
            val = self.right.pop(0)
            self.middle.add(val)
            self.middle_sum += val

    def calculateMKAverage(self) -> int:
        if len(self.stream) < self.m:
            return -1
        return self.middle_sum // len(self.middle)

TimeO(log m) per operation

SpaceO(m)

Open on LeetCode ↗