Problem statement

An unlimited row of stacks is numbered 0, 1, 2, … from left to right. Every stack holds at most capacity plates.

Write the class DinnerPlates:

• DinnerPlates(int capacity) creates the row, every stack empty.
• void push(int value) puts a plate with value on top of the leftmost stack that is not full.
• int pop() takes the top plate off the rightmost non-empty stack and returns its value, or returns -1 if every stack is empty.
• int popAtStack(int index) takes the top plate off stack index and returns its value, or returns -1 if that stack is empty.

The judge calls the methods in the order listed in the first input line, with the arguments in the second; the expected output holds each call's result, null for the constructor and for every method that returns nothing.

Example 1

Input: ["DinnerPlates","push","push","push","push","popAtStack","push","push","pop","popAtStack","pop","popAtStack"]
[[3],[4],[5],[6],[7],[0],[8],[9],[],[1],[],[5]]
Output: [null,null,null,null,null,6,null,null,9,7,8,-1]
Explanation: Plates 4, 5, 6 fill stack 0 and 7 starts stack 1. Taking 6 from stack 0 leaves room there, so 8 goes back to stack 0 and 9 to stack 1. pop takes 9 (stack 1 is rightmost), popAtStack(1) takes 7, and the next pop finds 8 on stack 0. Stack 5 is empty.

Example 2

Input: ["DinnerPlates","push","push","push","popAtStack","push","pop","pop","pop","pop"]
[[1],[1],[2],[3],[1],[7],[],[],[],[]]
Output: [null,null,null,null,2,null,3,7,1,-1]
Explanation: With capacity 1 each plate gets its own stack. Emptying stack 1 makes it the leftmost free one, so 7 lands there; pop then takes 3, 7 and 1, in that order.

Constraints

• 1 <= capacity <= 2 * 10^4
• 1 <= value <= 2 * 10^4
• 0 <= index <= 10^5
• At most 2 * 10^5 calls are made in total to push, pop and popAtStack

Hints for Dinner Plate Stacks

1. Hint 1

   Store the stacks in a growable list. pop is popAtStack on the last stack — keep the list trimmed so that its last stack is never empty.

2. Hint 2

   The hard part is push: a popAtStack in the middle can open a gap far to the left. Scanning for the leftmost non-full stack is O(n).

3. Hint 3

   Keep the indices of the non-full stacks in a min-heap (or an ordered set). popAtStack(i) adds i; push uses the smallest index and removes it once that stack is full. With a heap, discard tops that point past the end of the trimmed list or at a stack that is full again.

How to solve Dinner Plate Stacks

Two things must be found fast: the leftmost non-full stack (for push) and the rightmost non-empty one (for pop). A min-heap or ordered set of non-full stack indices answers the first; trimming empty stacks off the end of the list answers the second.

Approach

1. Keep stacks (a list of stacks, whose last stack is never empty) and open, the indices of stacks in the list that are not full.
2. push(value): if open is non-empty, push onto stack min(open) and drop that index from open once the stack reaches capacity. Otherwise append a new stack holding value (and add its index to open if capacity > 1).
3. popAtStack(index): if index is past the end or the stack is empty, return -1. Otherwise pop its top, add index to open, then remove empty stacks from the end of the list (and their indices from open).
4. pop(): return -1 if the list is empty, else popAtStack(last index).
5. With a lazy min-heap instead of an ordered set, push first discards tops that are past the end of the list or point at a full stack, and a flag per index keeps each index in the heap at most once.

Why it works

open always holds exactly the non-full stacks inside the list: a stack becomes non-full only when created or popped (both add it) and stops being non-full only when push fills it (which removes it). Every stack beyond the list is empty, so if open is empty the leftmost non-full stack is the next new one. Trimming keeps the last stack in the list non-empty, so pop acts on the rightmost non-empty stack.

Complexity

• Time — O(log n) per call (amortised for the trimming and the lazy heap)
• Space — O(number of plates + number of stacks)

Pitfalls

• pop must skip empty stacks on the right — trim them, or a middle popAtStack leaves pop looking at an empty stack.
• After trimming, indices in the heap may point past the end of the list; they are not valid targets for push until the list grows back.
• popAtStack with an index far beyond the list returns -1, not an error.

Reference solution

Python

import heapq

class DinnerPlates:

    def __init__(self, capacity: int):
        self.cap = capacity
        self.stacks = []   # the last stack is never empty
        self.heap = []     # candidate non-full stack indices (checked when used)
        self.inHeap = set()

    def _offer(self, i):
        if i not in self.inHeap:
            self.inHeap.add(i)
            heapq.heappush(self.heap, i)

    def push(self, value: int) -> None:
        while self.heap and (self.heap[0] >= len(self.stacks) or len(self.stacks[self.heap[0]]) == self.cap):
            self.inHeap.discard(heapq.heappop(self.heap))
        if self.heap:
            i = self.heap[0]
            self.stacks[i].append(value)
            if len(self.stacks[i]) == self.cap:
                self.inHeap.discard(heapq.heappop(self.heap))
        else:
            self.stacks.append([value])
            if self.cap > 1:
                self._offer(len(self.stacks) - 1)

    def pop(self) -> int:
        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
        x = self.stacks[index].pop()
        self._offer(index)
        while self.stacks and not self.stacks[-1]:
            self.stacks.pop()
        return x

Java

class DinnerPlates {
    private int cap;
    private ArrayList<ArrayDeque<Integer>> stacks = new ArrayList<>(); // the last stack is never empty
    private TreeSet<Integer> open = new TreeSet<>(); // indices of non-full stacks in the list

    public DinnerPlates(int capacity) {
        cap = capacity;
    }

    public void push(int value) {
        if (!open.isEmpty()) {
            int i = open.first();
            ArrayDeque<Integer> s = stacks.get(i);
            s.push(value);
            if (s.size() == cap) open.remove(i);
        } else {
            ArrayDeque<Integer> s = new ArrayDeque<>();
            s.push(value);
            stacks.add(s);
            if (cap > 1) open.add(stacks.size() - 1);
        }
    }

    public int pop() {
        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;
        int x = stacks.get(index).pop();
        open.add(index);
        while (!stacks.isEmpty() && stacks.get(stacks.size() - 1).isEmpty()) {
            open.remove(stacks.size() - 1);
            stacks.remove(stacks.size() - 1);
        }
        return x;
    }
}

C++

class DinnerPlates {
    int cap;
    vector<vector<int>> stacks; // the last stack is never empty
    set<int> open;              // indices of non-full stacks in the list
public:
    DinnerPlates(int capacity) : cap(capacity) {}

    void push(int value) {
        if (!open.empty()) {
            int i = *open.begin();
            stacks[i].push_back(value);
            if ((int)stacks[i].size() == cap) open.erase(i);
        } else {
            stacks.push_back({value});
            if (cap > 1) open.insert(stacks.size() - 1);
        }
    }

    int pop() {
        if (stacks.empty()) return -1;
        return popAtStack(stacks.size() - 1);
    }

    int popAtStack(int index) {
        if (index >= (int)stacks.size() || stacks[index].empty()) return -1;
        int x = stacks[index].back();
        stacks[index].pop_back();
        open.insert(index);
        while (!stacks.empty() && stacks.back().empty()) {
            open.erase(stacks.size() - 1);
            stacks.pop_back();
        }
        return x;
    }
};

JavaScript

var DinnerPlates = function(capacity) {
    this.cap = capacity;
    this.stacks = [];   // the last stack is never empty
    this.heap = [];     // candidate non-full stack indices (checked when used)
    this.inHeap = [];
};

DinnerPlates.prototype.offer = function(i) {
    if (this.inHeap[i]) return;
    this.inHeap[i] = true;
    var h = this.heap, k = h.length;
    h.push(i);
    while (k > 0 && h[(k - 1) >> 1] > i) { h[k] = h[(k - 1) >> 1]; k = (k - 1) >> 1; }
    h[k] = i;
};

DinnerPlates.prototype.take = function() { // remove the heap's minimum
    var h = this.heap;
    this.inHeap[h[0]] = false;
    var last = h.pop();
    if (h.length === 0) return;
    var k = 0, n = h.length;
    for (;;) {
        var l = 2 * k + 1, r = l + 1, m = k, mv = last;
        if (l < n && h[l] < mv) { m = l; mv = h[l]; }
        if (r < n && h[r] < mv) { m = r; }
        if (m === k) break;
        h[k] = h[m]; k = m;
    }
    h[k] = last;
};

DinnerPlates.prototype.push = function(value) {
    while (this.heap.length && (this.heap[0] >= this.stacks.length || this.stacks[this.heap[0]].length === this.cap)) this.take();
    if (this.heap.length) {
        var s = this.stacks[this.heap[0]];
        s.push(value);
        if (s.length === this.cap) this.take();
    } else {
        this.stacks.push([value]);
        if (this.cap > 1) this.offer(this.stacks.length - 1);
    }
};

DinnerPlates.prototype.pop = function() {
    if (this.stacks.length === 0) return -1;
    return this.popAtStack(this.stacks.length - 1);
};

DinnerPlates.prototype.popAtStack = function(index) {
    if (index >= this.stacks.length || this.stacks[index].length === 0) return -1;
    var x = this.stacks[index].pop();
    this.offer(index);
    while (this.stacks.length && this.stacks[this.stacks.length - 1].length === 0) this.stacks.pop();
    return x;
};

Also on the editorial tab: C, C#, Go, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

More hash table problems

• Design Movie Rental System Hard
• Prefix and Suffix Search Hard
• LFU Cache Hard
• Vertical Order Traversal of a Binary Tree Hard
• Maximum Number of Ways to Partition an Array Hard
• Sum of Imbalance Numbers of All Subarrays Hard