LeetCode 1172 - Dinner Plate Stacks

Difficulty: 🔴 Hard
Topics: Hash Table, Stack, Design, Heap (Priority Queue)

Solution

Problem Understanding

The problem asks us to design a data structure that behaves like an infinite row of stacks, where every stack has the same fixed capacity.

Instead of working with a single stack, we maintain many stacks indexed from left to right starting at 0. The operations are constrained by specific rules:

When we push(val), we must place the value into the leftmost stack that is not full.

When we pop(), we must remove and return the top element from the rightmost non-empty stack.

When we popAtStack(index), we remove and return the top element from the specific stack at index.

If no valid element exists for a pop operation, we return -1.

The tricky part is that stacks can become partially empty in the middle after popAtStack. Future pushes are required to reuse those earlier stacks before creating new stacks on the right.

For example, if capacity is 2 and we have:

Stack 0: [1]
Stack 1: [3,4]
Stack 2: [5]

then a future push(20) must go into stack 0, because it is the leftmost stack with available space.

The constraints are large:

At most 2 * 10^5 operations

This immediately tells us that linear scans during every operation will be too slow. A naive implementation that repeatedly scans all stacks would degrade to quadratic behavior.

Several edge cases are important:

If all stacks are empty, pop() must return -1.

If popAtStack(index) refers to an empty or non-existent stack, it must return -1.

After popping from the rightmost stack, trailing empty stacks should effectively disappear, otherwise future pop() operations may waste time scanning many empty stacks.

A stack that becomes non-full again after a pop must become eligible for future pushes.

The challenge is therefore maintaining two things efficiently:

The leftmost stack with free space.
The rightmost stack with elements.

Approaches

Brute Force Approach

The simplest solution is to maintain a list of stacks and directly scan it whenever needed.

For push(val), we scan from left to right until we find a stack whose size is less than capacity. If none exists, we create a new stack.

For pop(), we scan from right to left until we find a non-empty stack.

For popAtStack(index), we directly pop from that stack if possible.

This approach is correct because it directly follows the rules of the problem. However, it is inefficient.

Suppose we repeatedly pop elements from middle stacks, creating holes. Every future push may need to scan many stacks before finding an available slot. Similarly, pop() may repeatedly scan many empty stacks at the end.

In the worst case, each operation becomes O(n), where n is the number of stacks. With up to 2 * 10^5 operations, this becomes too slow.

Optimal Approach

The key observation is that we repeatedly need two special positions:

The leftmost non-full stack.
The rightmost non-empty stack.

Instead of scanning every time, we can maintain these efficiently.

We use:

A dynamic array of stacks.
A min-heap storing indices of stacks that are not full.

The heap allows us to quickly retrieve the smallest index with available capacity.

For the rightmost non-empty stack, we maintain the stacks array itself and lazily remove empty stacks from the end whenever necessary.

The important trick is lazy cleanup.

A stack index may remain inside the heap even after that stack becomes full again. Instead of aggressively removing invalid heap entries, we simply discard them when they appear at the top of the heap.

This gives efficient operations:

push: O(log n)
pop: amortized O(1)
popAtStack: O(log n)

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n) per operation	O(n)	Scans stacks repeatedly
Optimal	O(log n) per operation	O(n)	Uses heap for leftmost available stack

Algorithm Walkthrough

Initialize three fields:

capacity, the maximum size of each stack.
stacks, a list containing all stack arrays.
available, a min-heap storing indices of stacks that are not full.

During push(val), first clean invalid heap entries.

Some indices in the heap may refer to:

stacks that no longer exist
stacks that are already full

While the heap top is invalid, remove it. 3. If the heap becomes empty after cleanup, there is no existing stack with free space.

Create a new stack at the end:

stacks.append([])

Push its index into the heap because it initially has space. 4. The smallest index in the heap is the leftmost available stack.

Push the value into that stack. 5. After insertion, if the stack becomes full, remove its index from the heap.

This ensures future pushes skip full stacks. 6. During popAtStack(index), first verify:

the index exists
the stack is not empty

If invalid, return -1. 7. Pop the top value from the specified stack. 8. If the stack had been full before popping, it now has available space.

Push its index into the heap. 9. After popping, repeatedly remove trailing empty stacks from the end of stacks.

This keeps the rightmost non-empty stack accessible efficiently. 10. During pop(), first remove trailing empty stacks from the end. 11. If no stacks remain, return -1. 12. Otherwise, perform popAtStack on the last stack.

Why it works

The core invariant is:

The heap always contains indices that may have available capacity.
The right side of stacks never contains unnecessary trailing empty stacks.

The heap guarantees we always push into the leftmost non-full stack because it is a min-heap ordered by index.

Removing trailing empty stacks guarantees the last stack in stacks is always the rightmost non-empty stack.

Together, these properties exactly match the problem requirements.

Python Solution

from typing import List
import heapq

class DinnerPlates:

    def __init__(self, capacity: int):
        self.capacity = capacity
        self.stacks: List[List[int]] = []
        self.available: List[int] = []

    def push(self, val: int) -> None:
        while self.available:
            index = self.available[0]

            if index >= len(self.stacks):
                heapq.heappop(self.available)
            elif len(self.stacks[index]) == self.capacity:
                heapq.heappop(self.available)
            else:
                break

        if not self.available:
            self.stacks.append([])
            heapq.heappush(self.available, len(self.stacks) - 1)

        index = self.available[0]
        self.stacks[index].append(val)

        if len(self.stacks[index]) == self.capacity:
            heapq.heappop(self.available)

    def pop(self) -> int:
        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 < 0 or index >= len(self.stacks):
            return -1

        if not self.stacks[index]:
            return -1

        was_full = len(self.stacks[index]) == self.capacity

        value = self.stacks[index].pop()

        if was_full:
            heapq.heappush(self.available, index)

        while self.stacks and not self.stacks[-1]:
            self.stacks.pop()

        return value

The implementation closely follows the algorithm.

The stacks list stores every stack explicitly. Each stack itself is just a Python list.

The available min-heap stores candidate stack indices that may still have room. Because heap entries can become stale over time, the push operation performs lazy cleanup before using the heap top.

The push method either reuses the leftmost non-full stack or creates a new one if necessary.

The pop method removes trailing empty stacks before accessing the rightmost non-empty stack.

The popAtStack method handles direct indexed popping while also restoring the stack index to the heap if that stack becomes non-full again.

The repeated cleanup of trailing empty stacks ensures efficient future pop() operations.

Go Solution

package main

import (
	"container/heap"
)

type MinHeap []int

func (h MinHeap) Len() int {
	return len(h)
}

func (h MinHeap) Less(i, j int) bool {
	return h[i] < h[j]
}

func (h MinHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *MinHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *MinHeap) Pop() interface{} {
	old := *h
	n := len(old)
	val := old[n-1]
	*h = old[:n-1]
	return val
}

type DinnerPlates struct {
	capacity int
	stacks   [][]int
	available MinHeap
}

func Constructor(capacity int) DinnerPlates {
	h := MinHeap{}
	heap.Init(&h)

	return DinnerPlates{
		capacity: capacity,
		stacks:   [][]int{},
		available: h,
	}
}

func (this *DinnerPlates) Push(val int) {
	for this.available.Len() > 0 {
		index := this.available[0]

		if index >= len(this.stacks) {
			heap.Pop(&this.available)
		} else if len(this.stacks[index]) == this.capacity {
			heap.Pop(&this.available)
		} else {
			break
		}
	}

	if this.available.Len() == 0 {
		this.stacks = append(this.stacks, []int{})
		heap.Push(&this.available, len(this.stacks)-1)
	}

	index := this.available[0]
	this.stacks[index] = append(this.stacks[index], val)

	if len(this.stacks[index]) == this.capacity {
		heap.Pop(&this.available)
	}
}

func (this *DinnerPlates) Pop() int {
	for len(this.stacks) > 0 && len(this.stacks[len(this.stacks)-1]) == 0 {
		this.stacks = this.stacks[:len(this.stacks)-1]
	}

	if len(this.stacks) == 0 {
		return -1
	}

	return this.PopAtStack(len(this.stacks) - 1)
}

func (this *DinnerPlates) PopAtStack(index int) int {
	if index < 0 || index >= len(this.stacks) {
		return -1
	}

	if len(this.stacks[index]) == 0 {
		return -1
	}

	wasFull := len(this.stacks[index]) == this.capacity

	stack := this.stacks[index]
	value := stack[len(stack)-1]

	this.stacks[index] = stack[:len(stack)-1]

	if wasFull {
		heap.Push(&this.available, index)
	}

	for len(this.stacks) > 0 &&
		len(this.stacks[len(this.stacks)-1]) == 0 {
		this.stacks = this.stacks[:len(this.stacks)-1]
	}

	return value
}

The Go version follows the same logic as the Python implementation, but Go requires an explicit heap implementation using the container/heap package.

Stacks are represented using slices of slices:

[][]int

Unlike Python's built-in heapq, Go requires implementing heap interface methods such as Len, Less, Swap, Push, and Pop.

Slice operations are used heavily for stack push and pop behavior.

Worked Examples

Example 1

Input:

capacity = 2

Initial state:

Operation	Stacks	Available Heap
Initialize	[]	[]

push(1)

No available stack exists, create stack 0.

Operation	Stacks	Available Heap
push(1)	[[1]]	[0]

push(2)

Push into stack 0.

Now stack 0 becomes full.

Operation	Stacks	Available Heap
push(2)	[[1,2]]	[]

push(3)

Create stack 1.

Operation	Stacks	Available Heap
push(3)	[[1,2],[3]]	[1]

push(4)

Push into stack 1.

Operation	Stacks	Available Heap
push(4)	[[1,2],[3,4]]	[]

push(5)

Create stack 2.

Operation	Stacks	Available Heap
push(5)	[[1,2],[3,4],[5]]	[2]

popAtStack(0)

Pop value 2.

Stack 0 now has space again.

Operation	Stacks	Available Heap
popAtStack(0)	[[1],[3,4],[5]]	[0,2]

push(20)

Leftmost available stack is stack 0.

Operation	Stacks	Available Heap
push(20)	[[1,20],[3,4],[5]]	[2]

push(21)

Push into stack 2.

Operation	Stacks	Available Heap
push(21)	[[1,20],[3,4],[5,21]]	[]

popAtStack(0)

Pop value 20.

Operation	Stacks	Available Heap
popAtStack(0)	[[1],[3,4],[5,21]]	[0]

popAtStack(2)

Pop value 21.

Operation	Stacks	Available Heap
popAtStack(2)	[[1],[3,4],[5]]	[0,2]

pop()

Pop from rightmost non-empty stack.

Returns 5.

Operation	Stacks	Available Heap
pop()	[[1],[3,4]]	[0,2]

pop()

Returns 4.

Operation	Stacks	Available Heap
pop()	[[1],[3]]	[0,1,2]

pop()

Returns 3.

Operation	Stacks	Available Heap
pop()	[[1]]	[0,1,2]

pop()

Returns 1.

Operation	Stacks	Available Heap
pop()	[]	[0,1,2]

pop()

All stacks are empty.

Returns -1.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(log n) amortized	Heap operations dominate
Space	O(n)	Stores all elements and heap indices

Each push or popAtStack may insert or remove from the heap, which costs O(log n).

The cleanup loops are amortized efficient because each stack index or empty stack is removed only once over the entire execution.

Therefore, total runtime across all operations remains efficient enough for the problem constraints.

Test Cases

# Provided example
d = DinnerPlates(2)

d.push(1)
d.push(2)
d.push(3)
d.push(4)
d.push(5)

assert d.popAtStack(0) == 2  # pop specific stack

d.push(20)
d.push(21)

assert d.popAtStack(0) == 20  # reused leftmost stack
assert d.popAtStack(2) == 21  # pop from later stack

assert d.pop() == 5  # rightmost non-empty
assert d.pop() == 4
assert d.pop() == 3
assert d.pop() == 1
assert d.pop() == -1  # all empty

# Capacity 1
d = DinnerPlates(1)

d.push(10)
d.push(20)

assert d.popAtStack(0) == 10  # single element stack
assert d.pop() == 20
assert d.pop() == -1

# Reusing holes
d = DinnerPlates(2)

d.push(1)
d.push(2)
d.push(3)

assert d.popAtStack(0) == 2

d.push(4)

assert d.popAtStack(0) == 4  # hole reused before new stack

# Pop from invalid index
d = DinnerPlates(2)

assert d.popAtStack(0) == -1  # nonexistent stack

d.push(1)

assert d.popAtStack(5) == -1  # out of range

# Multiple trailing empty stacks
d = DinnerPlates(2)

d.push(1)
d.push(2)
d.push(3)
d.push(4)

assert d.pop() == 4
assert d.pop() == 3
assert d.pop() == 2
assert d.pop() == 1
assert d.pop() == -1

# Stress alternating operations
d = DinnerPlates(3)

for i in range(10):
    d.push(i)

for _ in range(5):
    d.pop()

for i in range(100, 105):
    d.push(i)

assert d.pop() == 104

Test Summary

Test	Why
Provided example	Validates all operations together
Capacity 1	Smallest stack size
Reusing holes	Ensures leftmost available stack is reused
Invalid index	Confirms bounds handling
Multiple trailing empties	Validates cleanup logic
Stress alternating operations	Tests heap consistency under mixed operations

Edge Cases

One important edge case occurs when all stacks become empty after several pop operations. A naive implementation may leave many empty stacks at the end of the array, causing future pop() operations to repeatedly scan useless stacks. This implementation continuously removes trailing empty stacks, ensuring the rightmost stack is always meaningful.

Another subtle case occurs when popAtStack creates a hole in the middle. Future pushes must reuse that stack before creating a new stack on the right. The min-heap guarantees that the smallest available index is always chosen first.

A third tricky case involves stale heap entries. Suppose a stack index is added to the heap because it became non-full, but later becomes full again before that heap entry is removed. Without cleanup, future pushes could incorrectly target a full stack. The implementation solves this using lazy deletion, repeatedly discarding invalid heap entries before each push operation.