给你两个整数 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 <= 105
1 <= k*2 < m
1 <= num <= 105
addElement
与calculateMKAverage
总操作次数不超过105
次。
方法一:有序集合 + 队列
我们可以维护以下数据结构或变量:
- 一个长度为
$m$ 的队列$q$ ,其中队首元素为最早加入的元素,队尾元素为最近加入的元素; - 三个有序集合,分别为
$lo$ ,$mid$ ,$hi$ ,其中$lo$ 和$hi$ 分别存储最小的$k$ 个元素和最大的$k$ 个元素,而$mid$ 存储剩余的元素; - 一个变量
$s$ ,维护$mid$ 中所有元素的和; - 部分编程语言(如 Java, Go)额外维护两个变量
$size1$ 和$size3$ ,分别表示$lo$ 和$hi$ 中元素的个数。
调用
- 如果
$lo$ 为空,或者$num \leq max(lo)$ ,则将$num$ 加入$lo$ 中;否则如果$hi$ 为空,或者$num \geq min(hi)$ ,则将$num$ 加入$hi$ 中;否则将$num$ 加入$mid$ 中,同时将$num$ 的值加到$s$ 中。 - 接下来将
$num$ 加入队列$q$ 中,如果此时队列$q$ 的长度大于$m$ ,则将队首元素$x$ 从队列$q$ 中移除,接下来从$lo$ ,$mid$ 或$hi$ 中选择其中一个包含$x$ 的集合,将$x$ 从该集合中移除,如果该集合为$mid$ ,则将$s$ 减去$x$ 的值。 - 如果
$lo$ 的长度大于$k$ ,则循环将$lo$ 中的最大值$max(lo)$ 从$lo$ 中移除,将$max(lo)$ 加入$mid$ 中,同时将$s$ 加上$max(lo)$ 的值。 - 如果
$hi$ 的长度大于$k$ ,则循环将$hi$ 中的最小值$min(hi)$ 从$hi$ 中移除,将$min(hi)$ 加入$mid$ 中,同时将$s$ 加上$min(hi)$ 的值。 - 如果
$lo$ 的长度小于$k$ ,并且$mid$ 不为空,则循环将$mid$ 中的最小值$min(mid)$ 从$mid$ 中移除,将$min(mid)$ 加入$lo$ 中,同时将$s$ 减去$min(mid)$ 的值。 - 如果
$hi$ 的长度小于$k$ ,并且$mid$ 不为空,则循环将$mid$ 中的最大值$max(mid)$ 从$mid$ 中移除,将$max(mid)$ 加入$hi$ 中,同时将$s$ 减去$max(mid)$ 的值。
调用
时间复杂度方面,每次调用
from sortedcontainers import SortedList
class MKAverage:
def __init__(self, m: int, k: int):
self.m = m
self.k = k
self.s = 0
self.q = deque()
self.lo = SortedList()
self.mid = SortedList()
self.hi = SortedList()
def addElement(self, num: int) -> None:
if not self.lo or num <= self.lo[-1]:
self.lo.add(num)
elif not self.hi or num >= self.hi[0]:
self.hi.add(num)
else:
self.mid.add(num)
self.s += num
self.q.append(num)
if len(self.q) > self.m:
x = self.q.popleft()
if x in self.lo:
self.lo.remove(x)
elif x in self.hi:
self.hi.remove(x)
else:
self.mid.remove(x)
self.s -= x
while len(self.lo) > self.k:
x = self.lo.pop()
self.mid.add(x)
self.s += x
while len(self.hi) > self.k:
x = self.hi.pop(0)
self.mid.add(x)
self.s += x
while len(self.lo) < self.k and self.mid:
x = self.mid.pop(0)
self.lo.add(x)
self.s -= x
while len(self.hi) < self.k and self.mid:
x = self.mid.pop()
self.hi.add(x)
self.s -= x
def calculateMKAverage(self) -> int:
return -1 if len(self.q) < self.m else self.s // (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()
from sortedcontainers import SortedList
class MKAverage:
def __init__(self, m: int, k: int):
self.m = m
self.k = k
self.sl = SortedList()
self.q = deque()
self.s = 0
def addElement(self, num: int) -> None:
self.q.append(num)
if len(self.q) == self.m:
self.sl = SortedList(self.q)
self.s = sum(self.sl[self.k: -self.k])
elif len(self.q) > self.m:
i = self.sl.bisect_left(num)
if i < self.k:
self.s += self.sl[self.k - 1]
elif self.k <= i <= self.m - self.k:
self.s += num
else:
self.s += self.sl[self.m - self.k]
self.sl.add(num)
x = self.q.popleft()
i = self.sl.bisect_left(x)
if i < self.k:
self.s -= self.sl[self.k]
elif self.k <= i <= self.m - self.k:
self.s -= x
else:
self.s -= self.sl[self.m - self.k]
self.sl.remove(x)
def calculateMKAverage(self) -> int:
return -1 if len(self.sl) < self.m else self.s // (self.m - self.k * 2)
# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()
class MKAverage {
private int m, k;
private long s;
private int size1, size3;
private Deque<Integer> q = new ArrayDeque<>();
private TreeMap<Integer, Integer> lo = new TreeMap<>();
private TreeMap<Integer, Integer> mid = new TreeMap<>();
private TreeMap<Integer, Integer> hi = new TreeMap<>();
public MKAverage(int m, int k) {
this.m = m;
this.k = k;
}
public void addElement(int num) {
if (lo.isEmpty() || num <= lo.lastKey()) {
lo.merge(num, 1, Integer::sum);
++size1;
} else if (hi.isEmpty() || num >= hi.firstKey()) {
hi.merge(num, 1, Integer::sum);
++size3;
} else {
mid.merge(num, 1, Integer::sum);
s += num;
}
q.offer(num);
if (q.size() > m) {
int x = q.poll();
if (lo.containsKey(x)) {
if (lo.merge(x, -1, Integer::sum) == 0) {
lo.remove(x);
}
--size1;
} else if (hi.containsKey(x)) {
if (hi.merge(x, -1, Integer::sum) == 0) {
hi.remove(x);
}
--size3;
} else {
if (mid.merge(x, -1, Integer::sum) == 0) {
mid.remove(x);
}
s -= x;
}
}
for (; size1 > k; --size1) {
int x = lo.lastKey();
if (lo.merge(x, -1, Integer::sum) == 0) {
lo.remove(x);
}
mid.merge(x, 1, Integer::sum);
s += x;
}
for (; size3 > k; --size3) {
int x = hi.firstKey();
if (hi.merge(x, -1, Integer::sum) == 0) {
hi.remove(x);
}
mid.merge(x, 1, Integer::sum);
s += x;
}
for (; size1 < k && !mid.isEmpty(); ++size1) {
int x = mid.firstKey();
if (mid.merge(x, -1, Integer::sum) == 0) {
mid.remove(x);
}
s -= x;
lo.merge(x, 1, Integer::sum);
}
for (; size3 < k && !mid.isEmpty(); ++size3) {
int x = mid.lastKey();
if (mid.merge(x, -1, Integer::sum) == 0) {
mid.remove(x);
}
s -= x;
hi.merge(x, 1, Integer::sum);
}
}
public int calculateMKAverage() {
return q.size() < m ? -1 : (int) (s / (q.size() - k * 2));
}
}
/**
* Your MKAverage object will be instantiated and called as such:
* MKAverage obj = new MKAverage(m, k);
* obj.addElement(num);
* int param_2 = obj.calculateMKAverage();
*/
class MKAverage {
public:
MKAverage(int m, int k) {
this->m = m;
this->k = k;
}
void addElement(int num) {
if (lo.empty() || num <= *lo.rbegin()) {
lo.insert(num);
} else if (hi.empty() || num >= *hi.begin()) {
hi.insert(num);
} else {
mid.insert(num);
s += num;
}
q.push(num);
if (q.size() > m) {
int x = q.front();
q.pop();
if (lo.find(x) != lo.end()) {
lo.erase(lo.find(x));
} else if (hi.find(x) != hi.end()) {
hi.erase(hi.find(x));
} else {
mid.erase(mid.find(x));
s -= x;
}
}
while (lo.size() > k) {
int x = *lo.rbegin();
lo.erase(prev(lo.end()));
mid.insert(x);
s += x;
}
while (hi.size() > k) {
int x = *hi.begin();
hi.erase(hi.begin());
mid.insert(x);
s += x;
}
while (lo.size() < k && mid.size()) {
int x = *mid.begin();
mid.erase(mid.begin());
s -= x;
lo.insert(x);
}
while (hi.size() < k && mid.size()) {
int x = *mid.rbegin();
mid.erase(prev(mid.end()));
s -= x;
hi.insert(x);
}
}
int calculateMKAverage() {
return q.size() < m ? -1 : s / (q.size() - k * 2);
}
private:
int m, k;
long long s = 0;
queue<int> q;
multiset<int> lo, mid, hi;
};
/**
* Your MKAverage object will be instantiated and called as such:
* MKAverage* obj = new MKAverage(m, k);
* obj->addElement(num);
* int param_2 = obj->calculateMKAverage();
*/
type MKAverage struct {
lo, mid, hi *redblacktree.Tree
q []int
m, k, s int
size1, size3 int
}
func Constructor(m int, k int) MKAverage {
lo := redblacktree.NewWithIntComparator()
mid := redblacktree.NewWithIntComparator()
hi := redblacktree.NewWithIntComparator()
return MKAverage{lo, mid, hi, []int{}, m, k, 0, 0, 0}
}
func (this *MKAverage) AddElement(num int) {
merge := func(rbt *redblacktree.Tree, key, value int) {
if v, ok := rbt.Get(key); ok {
nxt := v.(int) + value
if nxt == 0 {
rbt.Remove(key)
} else {
rbt.Put(key, nxt)
}
} else {
rbt.Put(key, value)
}
}
if this.lo.Empty() || num <= this.lo.Right().Key.(int) {
merge(this.lo, num, 1)
this.size1++
} else if this.hi.Empty() || num >= this.hi.Left().Key.(int) {
merge(this.hi, num, 1)
this.size3++
} else {
merge(this.mid, num, 1)
this.s += num
}
this.q = append(this.q, num)
if len(this.q) > this.m {
x := this.q[0]
this.q = this.q[1:]
if _, ok := this.lo.Get(x); ok {
merge(this.lo, x, -1)
this.size1--
} else if _, ok := this.hi.Get(x); ok {
merge(this.hi, x, -1)
this.size3--
} else {
merge(this.mid, x, -1)
this.s -= x
}
}
for ; this.size1 > this.k; this.size1-- {
x := this.lo.Right().Key.(int)
merge(this.lo, x, -1)
merge(this.mid, x, 1)
this.s += x
}
for ; this.size3 > this.k; this.size3-- {
x := this.hi.Left().Key.(int)
merge(this.hi, x, -1)
merge(this.mid, x, 1)
this.s += x
}
for ; this.size1 < this.k && !this.mid.Empty(); this.size1++ {
x := this.mid.Left().Key.(int)
merge(this.mid, x, -1)
this.s -= x
merge(this.lo, x, 1)
}
for ; this.size3 < this.k && !this.mid.Empty(); this.size3++ {
x := this.mid.Right().Key.(int)
merge(this.mid, x, -1)
this.s -= x
merge(this.hi, x, 1)
}
}
func (this *MKAverage) CalculateMKAverage() int {
if len(this.q) < this.m {
return -1
}
return this.s / (this.m - 2*this.k)
}
/**
* Your MKAverage object will be instantiated and called as such:
* obj := Constructor(m, k);
* obj.AddElement(num);
* param_2 := obj.CalculateMKAverage();
*/