Dinner Plate Stacks LeetCode solution in C++

1
class DinnerPlates {
2
 public:
3
  map<int, stack<int>> mp;
4
  set<int> empty;
5
  int cap;
6
  DinnerPlates(int capacity) {
7
    this->cap = capacity;
8
  }
9

10
  void push(int val) {
11
    if (empty.size() == 0) {
12
      empty.insert(mp.size());
13
    }
14
    mp[*empty.begin()].push(val);
15
    if (mp[*empty.begin()].size() == cap) {
16
      empty.erase(empty.begin());
17
    }
18
  }
19
  int pop() {
20
    if (mp.size() == 0) {
21
      return -1;
22
    }
23
    int index = mp.rbegin()->first;
24
    int val = mp[index].top();
25
    mp[index].pop();
26
    empty.insert(index);
27
    if (mp[index].size() == 0) {
28
      mp.erase(index);
29
    }
30
    return val;
31
  }
32

33
  int popAtStack(int index) {
34
    if (mp.size() == 0 || mp.find(index) == mp.end()) {
35
      return -1;
36
    }
37
    int val = mp[index].top();
38
    mp[index].pop();
39
    empty.insert(index);
40
    if (mp[index].size() == 0) {
41
      mp.erase(index);
42
    }
43
    return val;
44
  }
45
};