Dinner Plate Stacks
Problem
You have an infinite number of stacks arranged in a row, each with the same maximum capacity. Implement push(val) which pushes onto the leftmost stack that is not full, pop() which pops from the rightmost non-empty stack, and popAtStack(index) which pops from the stack at a given index. Any operation on an empty/non-existent stack returns -1.
capacity = 2; push(1) push(2) push(3) push(4) push(5); popAtStack(0); push(20) push(21); popAtStack(0); popAtStack(2); pop() pop() pop() pop() pop()2, 20, 21, 5, 4, 3, 1, -1import heapq
class DinnerPlates:
def __init__(self, capacity):
self.cap = capacity
self.stacks = []
self.avail = [] # min-heap of non-full stack indices
def push(self, val):
while self.avail and (self.avail[0] >= len(self.stacks)
or len(self.stacks[self.avail[0]]) == self.cap):
heapq.heappop(self.avail)
if not self.avail:
self.stacks.append([])
heapq.heappush(self.avail, len(self.stacks) - 1)
i = self.avail[0]
self.stacks[i].append(val)
if len(self.stacks[i]) == self.cap:
heapq.heappop(self.avail)
def pop(self):
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):
if index < 0 or index >= len(self.stacks) or not self.stacks[index]:
return -1
heapq.heappush(self.avail, index)
return self.stacks[index].pop()class MinHeap {
constructor() { this.a = []; }
size() { return this.a.length; }
peek() { return this.a[0]; }
push(x) {
const a = this.a; a.push(x); let i = a.length - 1;
while (i > 0) { const p = (i - 1) >> 1; if (a[p] <= a[i]) break; [a[p], a[i]] = [a[i], a[p]]; i = p; }
}
pop() {
const a = this.a, top = a[0], last = a.pop();
if (a.length) { a[0] = last; let i = 0;
for (;;) { let l = 2*i+1, r = 2*i+2, m = i;
if (l < a.length && a[l] < a[m]) m = l;
if (r < a.length && a[r] < a[m]) m = r;
if (m === i) break; [a[m], a[i]] = [a[i], a[m]]; i = m; } }
return top;
}
}
class DinnerPlates {
constructor(capacity) { this.cap = capacity; this.stacks = []; this.avail = new MinHeap(); }
push(val) {
const av = this.avail, st = this.stacks;
while (av.size() && (av.peek() >= st.length || st[av.peek()].length === this.cap)) av.pop();
if (!av.size()) { st.push([]); av.push(st.length - 1); }
const i = av.peek();
st[i].push(val);
if (st[i].length === this.cap) av.pop();
}
pop() {
const st = this.stacks;
while (st.length && st[st.length - 1].length === 0) st.pop();
if (!st.length) return -1;
return this.popAtStack(st.length - 1);
}
popAtStack(index) {
const st = this.stacks;
if (index < 0 || index >= st.length || st[index].length === 0) return -1;
this.avail.push(index);
return st[index].pop();
}
}class DinnerPlates {
int cap;
List<Deque<Integer>> stacks = new ArrayList<>();
PriorityQueue<Integer> avail = new PriorityQueue<>();
public DinnerPlates(int capacity) { cap = capacity; }
public void push(int val) {
while (!avail.isEmpty() && (avail.peek() >= stacks.size()
|| stacks.get(avail.peek()).size() == cap))
avail.poll();
if (avail.isEmpty()) {
stacks.add(new ArrayDeque<>());
avail.offer(stacks.size() - 1);
}
int i = avail.peek();
stacks.get(i).push(val);
if (stacks.get(i).size() == cap) avail.poll();
}
public int pop() {
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 < 0 || index >= stacks.size() || stacks.get(index).isEmpty())
return -1;
avail.offer(index);
return stacks.get(index).pop();
}
}class DinnerPlates {
int cap;
vector<vector<int>> stacks;
priority_queue<int, vector<int>, greater<int>> avail;
public:
DinnerPlates(int capacity) : cap(capacity) {}
void push(int val) {
while (!avail.empty() && (avail.top() >= (int)stacks.size()
|| (int)stacks[avail.top()].size() == cap))
avail.pop();
if (avail.empty()) {
stacks.push_back({});
avail.push((int)stacks.size() - 1);
}
int i = avail.top();
stacks[i].push_back(val);
if ((int)stacks[i].size() == cap) avail.pop();
}
int pop() {
while (!stacks.empty() && stacks.back().empty()) stacks.pop_back();
if (stacks.empty()) return -1;
return popAtStack((int)stacks.size() - 1);
}
int popAtStack(int index) {
if (index < 0 || index >= (int)stacks.size() || stacks[index].empty())
return -1;
avail.push(index);
int v = stacks[index].back();
stacks[index].pop_back();
return v;
}
};Explanation
We have a row of fixed-capacity stacks. The hard part is that push must always target the leftmost non-full stack, even after a popAtStack opens a gap somewhere in the middle. Scanning from the left every time would be slow, so we keep a min-heap of the indices of stacks that have free room.
The heap always hands back the smallest available index, which is exactly the leftmost stack we can push into. We store the stacks in a list and the heap in avail.
push(val) first cleans the heap top: it discards any index that is stale (past the end) or already full. If nothing is available, it appends a brand-new stack and adds its index. Then it pushes val onto stack i = avail[0] and, if that stack just became full, removes i from the heap.
pop() trims trailing empty stacks, then pops the rightmost non-empty stack. popAtStack(index) pops a specific stack and, crucially, pushes that index back into the heap because it now has room again.
Example with capacity 2: pushing 1..5 fills stacks [1,2] [3,4] [5]. popAtStack(0) removes 2 and frees index 0, so the next push(20) lands back in stack 0, and push(21) fills it again. Later pop calls drain from the right, returning 5, 4, 3, ... until everything is empty and further pops give -1.