Skip to content

1825. Finding MK Average

Difficulty Topics

Description

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:

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

Solution

finding-mk-average.py
from sortedcontainers import SortedList

class SortedListWithSum:
    def __init__(self):
        self.items = SortedList()
        self.total = 0

    def add(self, item):
        self.total += item
        self.items.add(item)

    def remove(self, item):
        self.total -= item
        self.items.remove(item)

class MKAverage:

    def __init__(self, m: int, k: int):
        # inside
        self.m = m
        # smallest
        self.k = k

        self.buffer = collections.deque()
        self.left = SortedListWithSum()
        self.mid = SortedListWithSum()
        self.right = SortedListWithSum()

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

        if len(self.buffer) > self.m:
            r = self.buffer.popleft()

            if r in self.right.items:
                self.right.remove(r)
            elif r in self.mid.items:
                self.mid.remove(r)

                small = self.right.items[0]
                self.right.remove(small)
                self.mid.add(small)
            else:
                self.left.remove(r)

                small = self.mid.items[0]
                self.mid.remove(small)
                self.left.add(small)

                small = self.right.items[0]
                self.right.remove(small)
                self.mid.add(small)

        self.left.add(num)
        if len(self.left.items) > self.k:
            big = self.left.items[-1]
            self.left.remove(big)

            self.mid.add(big)
            if len(self.mid.items) > self.m - self.k - self.k:
                big = self.mid.items[-1]
                self.mid.remove(big)
                self.right.add(big)
        #print(self.left.items, self.mid.items,  self.right.items, self.mid.total, self.m, self.k)        

    def calculateMKAverage(self) -> int:
        if len(self.buffer) >= self.m:
            #print(self.left.items, self.mid.items,  self.right.items, self.mid.total, self.m, self.k)
            return int(self.mid.total / (self.m - self.k - self.k))
        else:
            return -1




# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()