Skip to content

1172. Dinner Plate Stacks 👍

  • Time: Constructor: $O(1)$, push(val: int): $O(\log n)$, pop(): $O(\log n)$, popAtStack(index: int): $O(\log n)$
  • Space: $O(|\texttt{push()}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class DinnerPlates {
 public:
  DinnerPlates(int capacity) : capacity(capacity) {
    minHeap.push(0);
  }

  void push(int val) {
    const int index = minHeap.top();
    // add new stack on demand
    if (index == stacks.size())
      stacks.push_back({});
    // push the new value
    stacks[index].push(val);
    // if the stack pushed is full, remove its candidacy from minHeap
    if (stacks[index].size() == capacity) {
      minHeap.pop();
      // if minHeap is empty, the next available stack index is stacks.size()
      if (minHeap.empty())
        minHeap.push(stacks.size());
    }
  }

  int pop() {
    // remove empty stacks from back
    while (!stacks.empty() && stacks.back().empty())
      stacks.pop_back();
    if (stacks.empty())
      return -1;
    return popAtStack(stacks.size() - 1);
  }

  int popAtStack(int index) {
    if (index >= stacks.size() || stacks[index].empty())
      return -1;
    // if the stack is going to have space, add its candiday to minHeap
    if (stacks[index].size() == capacity)
      minHeap.push(index);
    const int val = stacks[index].top();
    stacks[index].pop();
    return val;
  }

 private:
  const int capacity;
  vector<stack<int>> stacks;
  // min stacks to push
  priority_queue<int, vector<int>, greater<>> minHeap;
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class DinnerPlates {
  public DinnerPlates(int capacity) {
    this.capacity = capacity;
    minHeap.offer(0);
  }

  public void push(int val) {
    final int index = minHeap.peek();
    // add new stack on demand
    if (index == stacks.size())
      stacks.add(new ArrayDeque<>());
    // push the new value
    stacks.get(index).push(val);
    // if the stack pushed is full, remove its candidacy from minHeap
    if (stacks.get(index).size() == capacity) {
      minHeap.poll();
      // if minHeap is empty, the next available stack index is stacks.size()
      if (minHeap.isEmpty())
        minHeap.offer(stacks.size());
    }
  }

  public int pop() {
    // remove empty stacks from back
    while (!stacks.isEmpty() && stacks.get(stacks.size() - 1).isEmpty())
      stacks.remove(stacks.size() - 1);
    if (stacks.isEmpty())
      return -1;
    return popAtStack(stacks.size() - 1);
  }

  public int popAtStack(int index) {
    if (index >= stacks.size() || stacks.get(index).isEmpty())
      return -1;
    // if the stack is going to have space, add its candiday to minHeap
    if (stacks.get(index).size() == capacity)
      minHeap.offer(index);
    return stacks.get(index).pop();
  }

  private final int capacity;
  private List<Deque<Integer>> stacks = new ArrayList<>();
  private Queue<Integer> minHeap = new PriorityQueue<>();
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class DinnerPlates:
  def __init__(self, capacity: int):
    self.capacity = capacity
    self.stacks = []
    self.minHeap = [0]  # min stacks to push

  def push(self, val: int) -> None:
    index = self.minHeap[0]
    # add new stack on demand
    if index == len(self.stacks):
      self.stacks.append([])
    # push the new value
    self.stacks[index].append(val)
    # if the stack pushed is full, remove its candidacy from minHeap
    if len(self.stacks[index]) == self.capacity:
      heapq.heappop(self.minHeap)
      # if minHeap is empty, the next available stack index is stacks.size()
      if not self.minHeap:
        heapq.heappush(self.minHeap, len(self.stacks))

  def pop(self) -> int:
    # remove empty stacks from back
    while self.stacks and not self.stacks[-1]:
      self.stacks.pop()
    if not self.stacks:
      return -1
    return self.popAtStack(len(self.stacks) - 1)

  def popAtStack(self, index: int) -> int:
    if index >= len(self.stacks) or not self.stacks[index]:
      return -1
    # if the stack is going to have space, add its candiday to minHeap
    if len(self.stacks[index]) == self.capacity:
      heapq.heappush(self.minHeap, index)
    return self.stacks[index].pop()