跳转至

1825. 求出 MK 平均值#

问题描述#

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

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

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

请你实现 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 <= 105
  • 1 <= k*2 < m
  • 1 <= num <= 105
  • addElement 与 calculateMKAverage 总操作次数不超过 105 次。

解题思路#

方法一：维护三个 multiset#

维护三个 multiset，分别是 left, mid, right，这三个 multiset 保存了最后 \(m\) 个元素，若将最后 \(m\) 个元素排序，left 保存了最小的 \(k\) 个数，mid 保存了中间的 \(m-2k\) 个数，right 保存了最大的 \(k\) 个数。

因为要将 left 的最大值移到 mid 中，所以 left 为降序排列，同理 mid 也降序排列，而 right 升序排列，即定义如下：

1
2
multiset<int, greater<>> left, mid;
multiset<int> right;

当在 mid 中进行元素的添加或者删除时，用一个变量 sum 对 mid 的元素之和进行记录。

当调用 addElement 时，用一个数组 history 记录添加的元素。

left, mid, right 元素的变化情况如下：

每次调用 addElement(num) 时，首先将 num 添加到 left 中。

如果 left.size() > k，则将 left 的最大值移到 mid 中。

如果 mid.size() > m - 2 * k，则将 mid 的最大值移到 right 中。

如果 right.size() > k，则说明需要在 left, mid, right 中删掉不在窗口中的数 x。

如果 x 在 right 中，则直接在 right 中删掉这个数就可以。

否则，把 right 中的最小值移到 mid 中，如果 x 在 mid 中，则直接在 mid 中删掉这个数。

否则，x 肯定在 left 中，把 mid 中的最小值移到 left 中，然后在 left 中删掉 x。

在上述过程中，记录 mid 中元素和的变化即可。

 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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
class MKAverage {
public:
    int m, k;
    long long sum;
    vector<int> history;
    multiset<int, greater<>> left, mid;
    multiset<int> right;

    MKAverage(int m, int k) : m(m), k(k), sum(0) {

    }

    void addElement(int num) {
        history.emplace_back(num);

        left.emplace(num);

        if (left.size() > k) {
            mid.emplace(*left.begin());
            sum += *left.begin();
            left.erase(left.begin());
        }

        if (mid.size() > m - 2 * k) {
            right.emplace(*mid.begin());
            sum -= *mid.begin();
            mid.erase(mid.begin());
        }

        if (right.size() > k) {
            int n = history.size();
            int x = history[n - m - 1]; // 不在窗口中的数
            auto it = right.find(x);

            if (it != right.end()) {
                right.erase(it);
            } else {
                mid.emplace(*right.begin());
                sum += *right.begin();
                right.erase(right.begin());

                it = mid.find(x);
                if (it != mid.end()) {
                    sum -= x;
                    mid.erase(it);
                } else {
                    it = prev(mid.end());
                    sum -= *it;
                    left.emplace(*it);
                    mid.erase(it);

                    it = left.find(x);
                    left.erase(it);
                }
            }
        }
    }

    int calculateMKAverage() {
        if ((int) right.size() != k) {
            return -1;
        }

        return sum / (m - 2 * k);
    }
};

方法二：Python 大法#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class MKAverage:

    def __init__(self, m: int, k: int):
        self.nums = []
        self.sorted_nums = []
        self.m = m
        self.k = k

    def addElement(self, num: int) -> None:
        self.nums.append(num)
        k = bisect_left(self.sorted_nums, num)
        self.sorted_nums[k:k] = [num]

        if len(self.sorted_nums) > self.m:
            k = bisect_left(self.sorted_nums, self.nums[-self.m - 1])
            self.sorted_nums[k : k + 1] = []
            # del self.sorted_nums[k]

    def calculateMKAverage(self) -> int:
        if len(self.sorted_nums) != self.m:
            return -1

        return sum(self.sorted_nums[self.k : self.m - self.k]) // (self.m - 2 * self.k)
返回顶部

在手机上阅读