1825：求出 MK 平均值（2395 分）

Ian

2021-04-11 约 802 字预计阅读 2 分钟次阅读

力扣第 236 场周赛第 4 题

题目

给你两个整数 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⁵ 次。

分析

考虑维护一个有序集合 sl，并动态维护 s=sum(sl[k:-k])。

弹出元素 x 时：

if x<sl[k]，弹出后 s[k] 补到前面， s -= sl[k]
elif i>=sl[-k]，弹出后 s[-k-1] 补到后面，s -= sl[-k-1]
else， s -= x

插入 num 时类似讨论：

if num<sl[k]，sl[k-1] 补到中间，s += slk-1]
elif num>sl[-k]，sl[-k] 补到中间，s += sl[-k]
else，s += num

解答

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38


class MKAverage:

    def __init__(self, m: int, k: int):
        from sortedcontainers import SortedList
        self.sl = SortedList()
        self.s = 0
        self.Q = deque()
        self.m = m
        self.k = k
        self.n = m-2*k

    def addElement(self, num: int) -> None:
        self.Q.append(num)
        if len(self.Q)<=self.m:
            self.sl.add(num)
            if len(self.Q)==self.m:
                self.s = sum(self.sl[self.k:-self.k])
            return 
        x = self.Q.popleft()
        if x<=self.sl[self.k-1]:
            self.s -= self.sl[self.k]
        elif x>=self.sl[-self.k]:
            self.s -= self.sl[-self.k-1]
        else:
            self.s -= x
        self.sl.remove(x)
        if num<self.sl[self.k-1]:
            self.s += self.sl[self.k-1]
        elif num>self.sl[-self.k]:
            self.s += self.sl[-self.k]
        else:
            self.s += num
        self.sl.add(num)

    def calculateMKAverage(self) -> int:
        if len(self.Q)<self.m:
            return -1
        return floor(self.s/self.n)

980 ms

目录

1825：求出 MK 平均值（2395 分）

题目

分析

解答