求出 MK 平均值

难度: Hard

给你两个整数 m 和 k ，以及数据流形式的若干整数。你需要实现一个数据结构，计算这个数据流的 MK 平均值 。

MK 平均值 按照如下步骤计算：

如果数据流中的整数少于 m 个，MK 平均值 为 -1 ，否则将数据流中最后 m 个元素拷贝到一个独立的容器中。
从这个容器中删除最小的 k 个数和最大的 k 个数。
计算剩余元素的平均值，并 向下取整到最近的整数 。

请你实现 MKAverage 类：

MKAverage(int m, int k) 用一个空的数据流和两个整数 m 和 k 初始化 MKAverage 对象。
void addElement(int num) 往数据流中插入一个新的元素 num 。
int calculateMKAverage() 对当前的数据流计算并返回 MK 平均数 ，结果需 向下取整到最近的整数 。

示例 1：

输入：
["MKAverage", "addElement", "addElement", "calculateMKAverage", "addElement", "calculateMKAverage", "addElement", "addElement", "addElement", "calculateMKAverage"]
[[3, 1], [3], [1], [], [10], [], [5], [5], [5], []]
输出：
[null, null, null, -1, null, 3, null, null, null, 5]

解释：
MKAverage obj = new MKAverage(3, 1); 
obj.addElement(3);        // 当前元素为 [3]
obj.addElement(1);        // 当前元素为 [3,1]
obj.calculateMKAverage(); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
obj.addElement(10);       // 当前元素为 [3,1,10]
obj.calculateMKAverage(); // 最后 3 个元素为 [3,1,10]
                          // 删除最小以及最大的 1 个元素后，容器为 [3]
                          // [3] 的平均值等于 3/1 = 3 ，故返回 3
obj.addElement(5);        // 当前元素为 [3,1,10,5]
obj.addElement(5);        // 当前元素为 [3,1,10,5,5]
obj.addElement(5);        // 当前元素为 [3,1,10,5,5,5]
obj.calculateMKAverage(); // 最后 3 个元素为 [5,5,5]
                          // 删除最小以及最大的 1 个元素后，容器为 [5]
                          // [5] 的平均值等于 5/1 = 5 ，故返回 5

提示：

3 <= m <= 10⁵
1 <= k*2 < m
1 <= num <= 10⁵
addElement 与 calculateMKAverage 总操作次数不超过 10⁵ 次。

Submission

运行时间: 541 ms

内存: 50.6 MB

from sortedcontainers import SortedList

class MKAverage:
    def __init__(self, m: int, k: int):
        self.m = m
        self.k = k
        self.queue = collections.deque([])  # last m elements
        self.lower = SortedList()   # smallest k
        self.middle = SortedList()  # middle elements: m-2*k element
        self.upper = SortedList()   # largest k
        self.sum = 0  # sum of the middle elements

    def addElement(self, num: int) -> None:
        self.queue.append(num)
        if len(self.lower) == 0 or num <= self.lower[-1]:
            self.lower.add(num)
        elif len(self.upper) == 0 or num >= self.upper[0]:
            self.upper.add(num)
        else:
            self.middle.add(num)
            self.sum += num
        
        if len(self.queue) > self.m:
            x = self.queue.popleft()
            if x in self.lower:
                self.lower.remove(x)
            elif x in self.upper:
                self.upper.remove(x)
            elif x in self.middle:
                self.middle.remove(x)
                self.sum -= x

        # always keep the 3 sorted arr balance: i.e. in their correct size
        while len(self.lower) > self.k:
            x = self.lower.pop()
            self.middle.add(x)
            self.sum += x
        while len(self.upper) > self.k:
            x = self.upper.pop(0)
            self.middle.add(x)
            self.sum += x
        while len(self.lower) < self.k and self.middle:
            x = self.middle.pop(0)
            self.lower.add(x)
            self.sum -= x
        while len(self.upper) < self.k and self.middle:
            x = self.middle.pop()
            self.upper.add(x)
            self.sum -= x



    def calculateMKAverage(self) -> int:
        if len(self.queue) < self.m:
            return -1
        return self.sum // (self.m-2*self.k)


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

Explain

该题解采用了三个有序列表（SortedList）和一个队列（deque）来维护数据流中的最后m个元素，并且实现了动态地维持这些元素中最小的k个、最大的k个和中间的m-2k个元素。每次添加新元素时，首先判断其应该插入哪个有序列表中，并且相应地更新这些列表和中间元素的和。当元素总数超过m时，从队列和相应的有序列表中移除最早的元素，并保持有序列表的平衡。计算MK平均值时，如果元素个数不足m，则返回-1；否则，返回中间元素和的整数除以中间元素的数量。

时间复杂度: O(log m)

空间复杂度: O(m)

from sortedcontainers import SortedList
import collections

class MKAverage:
    def __init__(self, m: int, k: int):
        self.m = m
        self.k = k
        self.queue = collections.deque([])  # 存储最新的m个元素
        self.lower = SortedList()   # 存储最小的k个数
        self.middle = SortedList()  # 存储中间的m-2*k个数
        self.upper = SortedList()   # 存储最大的k个数
        self.sum = 0  # 中间元素的和

    def addElement(self, num: int) -> None:
        self.queue.append(num)
        if len(self.lower) == 0 or num <= self.lower[-1]:
            self.lower.add(num)
        elif len(self.upper) == 0 or num >= self.upper[0]:
            self.upper.add(num)
        else:
            self.middle.add(num)
            self.sum += num
        
        if len(self.queue) > self.m:
            x = self.queue.popleft()
            if x in self.lower:
                self.lower.remove(x)
            elif x in self.upper:
                self.upper.remove(x)
            elif x in self.middle:
                self.middle.remove(x)
                self.sum -= x
        
        # 调整三个有序列表的大小，保持它们符合预期的长度
        while len(self.lower) > self.k:
            x = self.lower.pop()
            self.middle.add(x)
            self.sum += x
        while len(self.upper) > self.k:
            x = self.upper.pop(0)
            self.middle.add(x)
            self.sum += x
        while len(self.lower) < self.k and self.middle:
            x = self.middle.pop(0)
            self.lower.add(x)
            self.sum -= x
        while len(self.upper) < self.k and self.middle:
            x = self.middle.pop()
            self.upper.add(x)
            self.sum -= x

    def calculateMKAverage(self) -> int:
        if len(self.queue) < self.m:
            return -1
        return self.sum // (self.m-2*self.k)

Explore

保持三个有序列表的正确长度和元素顺序需要通过细心的管理元素添加和删除过程。首先，当元素被添加到系统中时，根据其与现有元素的比较，决定将其添加到最小值列表、最大值列表或中间值列表中。如果列表长度超过预定长度，将从一个列表的端点移动元素到另一个列表以确保所有列表保持正确的长度。例如，如果最小值列表超过了k个元素，将其最大的元素移动到中间值列表。此外，当元素从系统中删除（如从数据流中移除最老的元素）时，也需要检查并从相应的列表中移除该元素，之后再次调整列表以保持正确的长度。这个过程确保了每个列表都能在动态变化的情况下保持其应有的性质和顺序。

从中间列表移动元素到上层或下层列表时选择首元素或尾元素是为了维持列表的有序状态和逻辑一致性。具体来说，若需要向下层列表（存储较小元素）添加元素，应从中间列表的首部移动最小的元素；若需要向上层列表（存储较大元素）添加元素，则从中间列表的尾部移动最大的元素。这样做可以确保中间列表始终维持在正确的范围内，即不包括目前数据流中最大和最小的k个元素。这种选择直接影响到平均值的计算，因为它决定了哪些元素被纳入计算平均值的范围内，从而可能影响最终的平均值结果。

在删除元素时，首先确定该元素属于哪一个列表。这可以通过比较待删除元素与各列表中的元素范围进行判断。如果元素在最小值列表范围内，则从最小值列表中删除；如果在最大值列表范围内，则从最大值列表中删除；否则，从中间列表中删除。如果相同的元素值可能属于不同的列表，需要根据实际的元素位置和其在数据流中的状态来进行精确的管理。例如，如果一个值既可能是最小值列表中的最大值也可能是中间列表中的最小值，应基于元素在数据流中的顺序和列表的实际需求来决定从哪个列表中删除。这确保了即使在面对重复值时，列表的整体结构和计算逻辑也能保持正确无误。