Finding MK Average LeetCode solution in C++

1
/*
2
    Time: addElement: O(logm) | calculateMKAverage: O(1)
3
    Space: O(m)
4
    Tag: TreeMap, Sorting, Queue
5
    Difficulty: H
6
*/
7

8
class MKAverage {
9
  map<int, int> left, middle, right;
10
  queue<int> q;
11
  int sizeofLeft, sizeofMiddle, sizeofRight;
12
  int k;
13
  long long mkSum;
14
  int m;
15

16
  void addToSet1(int num) {
17
    left[num]++;
18
    sizeofLeft++;
19
  }
20

21
  void deleteFromSet1(int num) {
22
    left[num]--;
23
    if (left[num] == 0) left.erase(num);
24
    sizeofLeft--;
25
  }
26

27
  void addToSet2(int num) {
28
    middle[num]++;
29
    sizeofMiddle++;
30
    mkSum += num;
31
  }
32

33
  void deleteFromSet2(int num) {
34
    middle[num]--;
35
    if (middle[num] == 0) middle.erase(num);
36
    sizeofMiddle--;
37
    mkSum -= num;
38
  }
39

40
  void addToSet3(int num) {
41
    right[num]++;
42
    sizeofRight++;
43
  }
44

45
  void deleteFromSet3(int num) {
46
    right[num]--;
47
    if (right[num] == 0) right.erase(num);
48
    sizeofRight--;
49
  }
50

51
 public:
52
  MKAverage(int m, int k) {
53
    sizeofLeft = 0, sizeofMiddle = 0, sizeofRight = 0;
54
    mkSum = 0;
55
    this->k = k;
56
    this->m = m;
57
  }
58

59
  void addElement(int num) {
60
    if (sizeofLeft < k) {
61
      addToSet1(num);
62
      q.push(num);
63
    } else if (sizeofMiddle < (m - (2 * k))) {
64
      int lastEle = prev(left.end())->first;
65
      if (num >= lastEle) {
66
        addToSet2(num);
67
      } else {
68
        deleteFromSet1(lastEle);
69
        addToSet1(num);
70
        addToSet2(lastEle);
71
      }
72
      q.push(num);
73
    } else if (sizeofRight < k) {
74
      int last1 = prev(left.end())->first;
75
      int last2 = prev(middle.end())->first;
76
      if (num >= last1 && num >= last2) {
77
        addToSet3(num);
78
      } else if (num >= last1 && num < last2) {
79
        deleteFromSet2(last2);
80
        addToSet2(num);
81
        addToSet3(last2);
82
      } else {
83
        deleteFromSet2(last2);
84
        addToSet3(last2);
85
        deleteFromSet1(last1);
86
        addToSet2(last1);
87
        addToSet1(num);
88
      }
89
      q.push(num);
90
    } else {
91
      int toErase = q.front();
92
      q.pop();
93
      int first3 = right.begin()->first;
94
      int first2 = middle.begin()->first;
95
      int first1 = left.begin()->first;
96

97
      if (toErase >= first3) {
98
        deleteFromSet3(toErase);
99
      } else if (toErase >= first2) {
100
        deleteFromSet3(first3);
101
        deleteFromSet2(toErase);
102
        addToSet2(first3);
103
      } else {
104
        deleteFromSet3(first3);
105
        deleteFromSet2(first2);
106
        deleteFromSet1(toErase);
107
        addToSet1(first2);
108
        addToSet2(first3);
109
      }
110
      addElement(num);
111
    }
112
  }
113

114
  int calculateMKAverage() {
115
    if (q.size() < m) return -1;
116
    return mkSum / (m - (2 * k));
117
  }
118
};