feat: add solutions to lc problem: No.1825

No.1825.Finding MK Average

Author: yanglbme

solution/1800-1899/1825.Finding MK Average/README.md

@@ -67,22 +67,373 @@ obj.calculateMKAverage(); // 最后 3 个元素为 [5,5,5]
 
 <!-- 这里可写通用的实现逻辑 -->
 
+**方法一：有序集合 + 队列**
+
+在本题中，我们需要维护的数据结构或变量有：
+
+-   一个长度为 $m$ 的队列 $q$，其中队首元素为最早加入的元素，队尾元素为最近加入的元素
+-   三个有序集合，分别为 $lo$, $mid$, $hi$，其中 $lo$ 和 $hi$ 分别存储最小的 $k$ 个元素和最大的 $k$ 个元素，而 $mid$ 存储剩余的元素；
+-   一个变量 $s$，维护 $mid$ 中所有元素的和。
+-   部分编程语言（如 Java, Go）额外维护两个变量 $size1$ 和 $size3$，分别表示 $lo$ 和 $hi$ 中元素的个数。
+
+调用 $addElement(num)$ 函数时，顺序执行以下操作：
+
+1. 如果 $lo$ 为空，或者 $num \leq max(lo)$，则将 $num$ 加入 $lo$ 中；否则如果 $hi$ 为空，或者 $num \geq min(hi)$，则将 $num$ 加入 $hi$ 中；否则将 $num$ 加入 $mid$ 中，同时将 $num$ 的值加到 $s$ 中。
+1. 接下来将 $num$ 加入队列 $q$ 中，如果此时队列 $q$ 的长度大于 $m$，则将队首元素 $x$ 从队列 $q$ 中移除，接下来从 $lo$, $mid$ 或 $hi$ 中选择其中一个包含 $x$ 的集合，将 $x$ 从该集合中移除，如果该集合为 $mid$，则将 $s$ 减去 $x$ 的值。
+1. 如果 $lo$ 的长度大于 $k$，则循环将 $lo$ 中的最大值 $max(lo)$ 从 $lo$ 中移除，将 $max(lo)$ 加入 $mid$ 中，同时将 $s$ 加上 $max(lo)$ 的值。
+1. 如果 $hi$ 的长度大于 $k$，则循环将 $hi$ 中的最小值 $min(hi)$ 从 $hi$ 中移除，将 $min(hi)$ 加入 $mid$ 中，同时将 $s$ 加上 $min(hi)$ 的值。
+1. 如果 $lo$ 的长度小于 $k$，并且 $mid$ 不为空，则循环将 $mid$ 中的最小值 $min(mid)$ 从 $mid$ 中移除，将 $min(mid)$ 加入 $lo$ 中，同时将 $s$ 减去 $min(mid)$ 的值。
+1. 如果 $hi$ 的长度小于 $k$，并且 $mid$ 不为空，则循环将 $mid$ 中的最大值 $max(mid)$ 从 $mid$ 中移除，将 $max(mid)$ 加入 $hi$ 中，同时将 $s$ 减去 $max(mid)$ 的值。
+
+调用 $calculateMKAverage()$ 函数时，如果 $q$ 的长度小于 $m$，则返回 $-1$，否则返回 $\frac{s}{m - 2k}$。
+
+时间复杂度方面，每次调用 $addElement(num)$ 函数的时间复杂度为 $O(\log m)$，每次调用 $calculateMKAverage()$ 函数的时间复杂度为 $O(1)$。空间复杂度为 $O(m)$。
+
 <!-- tabs:start -->
 
 ### **Python3**
 
 <!-- 这里可写当前语言的特殊实现逻辑 -->
 
 ```python
+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()
 ```
 
 ### **Java**
 
 <!-- 这里可写当前语言的特殊实现逻辑 -->
 
 ```java
+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)) {
+                lo.merge(x, -1, Integer::sum);
+                if (lo.get(x) == 0) {
+                    lo.remove(x);
+                }
+                --size1;
+            } else if (hi.containsKey(x)) {
+                hi.merge(x, -1, Integer::sum);
+                if (hi.get(x) == 0) {
+                    hi.remove(x);
+                }
+                --size3;
+            } else {
+                mid.merge(x, -1, Integer::sum);
+                if (mid.get(x) == 0) {
+                    mid.remove(x);
+                }
+                s -= x;
+            }
+        }
+        for (; size1 > k; --size1) {
+            int x = lo.lastKey();
+            lo.merge(x, -1, Integer::sum);
+            if (lo.get(x) == 0) {
+                lo.remove(x);
+            }
+            mid.merge(x, 1, Integer::sum);
+            s += x;
+        }
+        for (; size3 > k; --size3) {
+            int x = hi.firstKey();
+            hi.merge(x, -1, Integer::sum);
+            if (hi.get(x) == 0) {
+                hi.remove(x);
+            }
+            mid.merge(x, 1, Integer::sum);
+            s += x;
+        }
+        for (; size1 < k && !mid.isEmpty(); ++size1) {
+            int x = mid.firstKey();
+            mid.merge(x, -1, Integer::sum);
+            if (mid.get(x) == 0) {
+                mid.remove(x);
+            }
+            s -= x;
+            lo.merge(x, 1, Integer::sum);
+        }
+        for (; size3 < k && !mid.isEmpty(); ++size3) {
+            int x = mid.lastKey();
+            mid.merge(x, -1, Integer::sum);
+            if (mid.get(x) == 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();
+ */
+```
+
+### **C++**
+
+```cpp
+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();
+ */
+```
+
+### **Go**
+
+```go
+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();
+ */
 ```
 
 ### **...**