1825. 找到 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^5
1 < k*2 < m
1 <= num <= 10^5
最多调用 10^5 次 addElement 和 calculateMKAverage

解题思路

这道题需要维护一个滑动窗口的数据结构，能够快速：

添加新元素并维护最近 m 个元素
快速找到并删除最小/最大的 k 个元素
计算中间元素的和

方法一：三个有序集合 将最近的 m 个元素分为三个部分：

small：存储最小的 k 个元素
middle：存储中间的 m-2k 个元素
large：存储最大的 k 个元素

使用队列维护元素顺序，用三个 multiset 分别维护这三部分，并记录 middle 部分的和。

方法二：单个有序集合 + 迭代器 使用一个 multiset 存储所有元素，通过迭代器快速定位边界，计算中间部分的和。

推荐使用方法一，因为它预先维护了中间部分的和，查询时间复杂度更优。

关键操作：

添加元素时，先加入对应的集合，然后重新平衡三个集合的大小
删除元素时，从对应集合中移除，然后重新平衡
平衡策略：确保 small 有 k 个元素，large 有 k 个元素，middle 有 m-2k 个元素

时间复杂度：addElement 和 calculateMKAverage 都是 O(log m)

代码实现

class MKAverage {
private:
    int m, k;
    queue<int> stream;
    multiset<int> small, middle, large;
    long long middleSum;
    
    void balance() {
        // Move elements to achieve correct sizes
        while (small.size() > k) {
            int val = *small.rbegin();
            small.erase(small.find(val));
            middle.insert(val);
            middleSum += val;
        }
        while (large.size() > k) {
            int val = *large.begin();
            large.erase(large.find(val));
            middle.insert(val);
            middleSum += val;
        }
        while (middle.size() > m - 2 * k) {
            if (small.size() < k) {
                int val = *middle.begin();
                middle.erase(middle.find(val));
                middleSum -= val;
                small.insert(val);
            } else {
                int val = *middle.rbegin();
                middle.erase(middle.find(val));
                middleSum -= val;
                large.insert(val);
            }
        }
        while (small.size() < k && !middle.empty()) {
            int val = *middle.rbegin();
            middle.erase(middle.find(val));
            middleSum -= val;
            small.insert(val);
        }
        while (large.size() < k && !middle.empty()) {
            int val = *middle.begin();
            middle.erase(middle.find(val));
            middleSum -= val;
            large.insert(val);
        }
    }
    
public:
    MKAverage(int m, int k) : m(m), k(k), middleSum(0) {}
    
    void addElement(int num) {
        stream.push(num);
        
        // Add to appropriate set
        if (small.size() < k) {
            small.insert(num);
        } else if (large.size() < k) {
            large.insert(num);
        } else {
            middle.insert(num);
            middleSum += num;
        }
        
        // Remove oldest element if needed
        if (stream.size() > m) {
            int old = stream.front();
            stream.pop();
            
            if (small.find(old) != small.end()) {
                small.erase(small.find(old));
            } else if (large.find(old) != large.end()) {
                large.erase(large.find(old));
            } else {
                middle.erase(middle.find(old));
                middleSum -= old;
            }
        }
        
        balance();
    }
    
    int calculateMKAverage() {
        if (stream.size() < m) return -1;
        return middleSum / (m - 2 * k);
    }
};

from collections import deque
from sortedcontainers import SortedList

class MKAverage:
    def __init__(self, m: int, k: int):
        self.m = m
        self.k = k
        self.stream = deque()
        self.small = SortedList()
        self.middle = SortedList()
        self.large = SortedList()
        self.middle_sum = 0
    
    def balance(self):
        # Move elements to achieve correct sizes
        while len(self.small) > self.k:
            val = self.small.pop()
            self.middle.add(val)
            self.middle_sum += val
        
        while len(self.large) > self.k:
            val = self.large.pop(0)
            self.middle.add(val)
            self.middle_sum += val
        
        while len(self.middle) > self.m - 2 * self.k:
            if len(self.small) < self.k:
                val = self.middle.pop(0)
                self.middle_sum -= val
                self.small.add(val)
            else:
                val = self.middle.pop()
                self.middle_sum -= val
                self.large.add(val)
        
        while len(self.small) < self.k and self.middle:
            val = self.middle.pop()
            self.middle_sum -= val
            self.small.add(val)
        
        while len(self.large) < self.k and self.middle:
            val = self.middle.pop(0)
            self.middle_sum -= val
            self.large.add(val)

    def addElement(self, num: int) -> None:
        self.stream.append(num)
        
        # Add to appropriate set
        if len(self.small) < self.k:
            self.small.add(num)
        elif len(self.large) < self.k:
            self.large.add(num)
        else:
            self.middle.add(num)
            self.middle_sum += num
        
        # Remove oldest element if needed
        if len(self.stream) > self.m:
            old = self.stream.popleft()
            
            if old in self.small:
                self.small.remove(old)
            elif old in self.large:
                self.large.remove(old)
            else:
                self.middle.remove(old)
                self.middle_sum -= old
        
        self.balance()

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

public class MKAverage {
    private int m, k;
    private Queue<int> stream;
    private SortedDictionary<int, int> small, middle, large;
    private long middleSum;
    private int smallCount, middleCount, largeCount;
    
    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
        stream = new Queue<int>();
        small = new SortedDictionary<int, int>();
        middle = new SortedDictionary<int, int>();
        large = new SortedDictionary<int, int>();
        middleSum = 0;
        smallCount = middleCount = largeCount = 0;
    }
    
    private void AddToSet(SortedDictionary<int, int> set, int val, ref int count) {
        if (!set.ContainsKey(val)) set[val] = 0;
        set[val]++;
        count++;
    }
    
    private void RemoveFromSet(SortedDictionary<int, int> set, int val, ref int count) {
        set[val]--;
        count--;
        if (set[val] == 0) set.Remove(val);
    }
    
    private void Balance() {
        // Move from small to middle
        while (smallCount > k) {
            var val = small.Keys.Last();
            RemoveFromSet(small, val, ref smallCount);
            AddToSet(middle, val, ref middleCount);
            middleSum += val;
        }
        
        // Move from large to middle
        while (largeCount > k) {
            var val = large.Keys.First();
            RemoveFromSet(large, val, ref largeCount);
            AddToSet(middle, val, ref middleCount);
            middleSum += val;
        }
        
        // Move from middle to small or large
        while (middleCount > m - 2 * k) {
            if (smallCount < k) {
                var val = middle.Keys.First();
                RemoveFromSet(middle, val, ref middleCount);
                middleSum -= val;
                AddToSet(small, val, ref smallCount);
            } else {
                var val = middle.Keys.Last();
                RemoveFromSet(middle, val, ref middleCount);
                middleSum -= val;
                AddToSet(large, val, ref largeCount);
            }
        }
        
        // Fill small from middle
        while (smallCount < k && middleCount > 0) {
            var val = middle.Keys.Last();
            RemoveFromSet(middle, val, ref middleCount);
            middleSum -= val;
            AddToSet(small, val, ref smallCount);
        }
        
        // Fill large from middle
        while (largeCount < k && middleCount > 0) {
            var val = middle.Keys.First();
            RemoveFromSet(middle, val, ref middleCount);
            middleSum -= val;
            AddToSet(large, val, ref largeCount);
        }
    }
    
    public void AddElement(int num) {
        stream.Enqueue(num);
        
        // Add to appropriate set
        if (smallCount < k) {
            AddToSet(small, num, ref smallCount);
        } else if (largeCount < k) {
            AddToSet(large, num, ref largeCount);
        } else {
            AddToSet(middle, num, ref middleCount);
            middleSum += num;
        }
        
        // Remove oldest element if needed
        if (stream.Count > m) {
            int old = stream.Dequeue();
            
            if (small.ContainsKey(old)) {
                RemoveFromSet(small, old, ref smallCount);
            } else if (large.ContainsKey(old)) {
                RemoveFromSet(large, old, ref largeCount);
            } else {
                RemoveFromSet(middle, old, ref middleCount);
                middleSum -= old;
            }
        }
        
        Balance();
    }
    
    public int CalculateMKAverage() {
        if (stream.Count < m) return -1;
        return (int)(middleSum / (m - 2 * k));
    }
}

var MKAverage = function(m, k) {
    this.m = m;
    this.k = k;
    this.stream = [];
};

MKAverage.prototype.addElement = function(num) {
    this.stream.push(num);
    if (this.stream.length > this.m) {
        this.stream.shift();
    }
};

MKAverage.prototype.calculateMKAverage = function() {
    if (this.stream.length < this.m) {
        return -1;
    }
    
    const sorted = [...this.stream].sort((a, b) => a - b);
    const middle = sorted.slice(this.k, this.m - this.k);
    const sum = middle.reduce((acc, val) => acc + val, 0);
    
    return Math.floor(sum / middle.length);
};

在 LeetCode 上运行：力扣中文版 LeetCode 国际版

复杂度分析

指标	复杂度
时间	-
空间	-

题目描述#

解题思路#

解题思路#

代码实现#

复杂度分析#

相关题目#

题目描述

解题思路

解题思路

代码实现

复杂度分析

相关题目