Dinner Plate Stacks LeetCode solution in Java

1
class DinnerPlates {
2

3
  List<Stack<Integer>> stack;
4
  PriorityQueue<Integer> leftEmpty;
5
  PriorityQueue<Integer> rightNonEmpty;
6
  int cap;
7

8
  public DinnerPlates(int capacity) {
9
    this.cap = capacity;
10
    this.stack = new ArrayList<>();
11
    this.leftEmpty = new PriorityQueue<>();
12
    this.rightNonEmpty = new PriorityQueue<>(Collections.reverseOrder());
13
    stack.add(new Stack<>());
14
    leftEmpty.offer(0);
15
  }
16

17
  public void push(int val) {
18
    while (!leftEmpty.isEmpty() && stack.get(leftEmpty.peek()).size() == cap) leftEmpty.poll();
19
    if (leftEmpty.isEmpty()) {
20
      stack.add(new Stack<>());
21
      leftEmpty.offer(stack.size() - 1);
22
    }
23
    Stack<Integer> s = stack.get(leftEmpty.peek());
24
    if (s.isEmpty()) rightNonEmpty.offer(leftEmpty.peek());
25
    s.push(val);
26
  }
27

28
  public int pop() {
29
    while (!rightNonEmpty.isEmpty() && stack.get(rightNonEmpty.peek()).isEmpty())
30
      rightNonEmpty.poll();
31
    if (rightNonEmpty.isEmpty()) {
32
      return -1;
33
    }
34
    Stack<Integer> s = stack.get(rightNonEmpty.peek());
35
    if (s.size() == cap) leftEmpty.offer(rightNonEmpty.peek());
36
    return s.pop();
37
  }
38

39
  public int popAtStack(int index) {
40
    if (index >= stack.size()) return -1;
41
    Stack<Integer> s = stack.get(index);
42
    if (s.isEmpty()) return -1;
43
    if (s.size() == cap) leftEmpty.offer(index);
44
    return s.pop();
45
  }
46
}
47

48
/**
49
 * Your DinnerPlates object will be instantiated and called as such: DinnerPlates obj = new
50
 * DinnerPlates(capacity); obj.push(val); int param_2 = obj.pop(); int param_3 =
51
 * obj.popAtStack(index);
52
 */