Editorial: LeetCode 1825 Finding Mk Average
https://leetcode.com/problems/finding-mk-average/ from Google
Thinking process:
The problem is why I prefer using C++ to do the coding interview. The advantage of using C++ is
- There are many competitive programmers using C++ as it’s faster than other languages, so you can find a lot of solutions from those good programmers
- There are a lot of good tools in C++'s library, like multiset in this problem, which is a ordered set which can store duplicate values
If you know multiset in C++, this problem is not very hard for you. We create three multiset first, top, middle, and bottom. Top is used for the greatest k elements. Bottom is used for the smallest k elements.
Whenever we add a new number, we add it to top first. Next, we push the smallest element of top to middle, and do the same thing to middle as well. In this way, we maintain an order in these three multisets.
Whenever we need to remove a number, we search the top first. If we remove a number from top, we need to pull a number from middle to top so that the size of top is still k. After this, we’ll need to pull another number from bottom to middle so that the size of middle is still m - 2*k. Same thing happens when we remove a number from middle.
Here is the code.
1class MKAverage {
2public:
3 typedef unsigned long long ll;
4 int M, K;
5 multiset<ll> top, middle, bottom;
6 ll sum;
7 queue<int> q;
8 MKAverage(int m, int k) {
9 M = m, K = k, sum = 0;
10 }
11
12 void addElement(int num) {
13 q.push(num);
14 add(num);
15 if (q.size() > M) {
16 int r = q.front(); q.pop();
17 remove(r);
18 }
19 }
20
21 int calculateMKAverage() {
22 if (q.size() < M) return -1;
23 return sum/(M-2*K);
24 }
25
26private:
27 void add(int num) {
28 top.insert(num);
29 if (top.size() > K) {
30 auto smallest = top.begin();
31 int s = *smallest;
32 top.erase(smallest);
33 middle.insert(s);
34 sum += s;
35 }
36 if (middle.size() > (M - 2*K)) {
37 auto smallest = middle.begin();
38 int s = *smallest;
39 middle.erase(smallest);
40 sum -= s;
41 bottom.insert(s);
42 }
43 }
44
45 void remove(int r) {
46 if (top.count(r)) {
47 auto it = top.find(r);
48 top.erase(it);
49 pullMiddle();
50 pullBottom();
51 } else if (middle.count(r)) {
52 auto it = middle.find(r);
53 sum -= *it;
54 middle.erase(it);
55 pullBottom();
56 } else {
57 auto it = bottom.find(r);
58 bottom.erase(it);
59 }
60 }
61
62 void pullMiddle() {
63 if (top.size() < K) {
64 auto it = middle.end();
65 it--;
66 top.insert(*it);
67 sum -= *it;
68 middle.erase(it);
69 }
70 }
71
72 void pullBottom() {
73 if (middle.size() < (M - 2*K)) {
74 auto it = bottom.end();
75 it--;
76 middle.insert(*it);
77 sum += *it;
78 bottom.erase(it);
79 }
80 }
81};
82
83/**
84 * Your MKAverage object will be instantiated and called as such:
85 * MKAverage* obj = new MKAverage(m, k);
86 * obj->addElement(num);
87 * int param_2 = obj->calculateMKAverage();
88 */