1172. Dinner Plate Stacks

Problem¶

Implement stacks of fixed capacity, indexed from left to right:

push(val): insert into the leftmost stack with room.
pop(): remove from the rightmost nonempty stack.
popAtStack(index): remove from that stack.
Both removal methods return -1 when no plate is available.

Stacks must keep their indices, even when empty.

capacity = 2                         stacks, bottom -> top
push(1), push(2), push(3)             [[1, 2], [3]]
popAtStack(0) -> 2                    [[1], [3]]
push(4)                              [[1, 4], [3]]
pop() -> 3                           [[1, 4], []]
popAtStack(1) -> -1                   [[1, 4], []]

class DinnerPlates:
    def __init__(self, capacity: int):
        pass

    def push(self, val: int) -> None:
        pass

    def pop(self) -> int:
        pass

    def popAtStack(self, index: int) -> int:
        pass

Key trick¶

Use a min-heap of nonfull stack indices and a pointer to the rightmost nonempty stack; add an index to the heap only when its stack transitions from full to nonfull.

Trap¶

Duplicate heap entries can cause overfilling; removing interior stacks changes indices, and removing trailing stacks requires handling stale heap entries.

Why is it interesting?¶

It combines two different ordering requirements: the smallest available index for insertion and the largest occupied index for removal.

Python solution¶

class DinnerPlates:
    def __init__(self, capacity: int):
        self.capacity = capacity
        self.stacks = []
        self.heap = []
        self.right = -1

    def push(self, val: int) -> None:
        if not self.heap:
            # Retain empty stacks so indices and heap entries stay valid.
            heapq.heappush(self.heap, len(self.stacks))
            self.stacks.append([])

        i = self.heap[0]
        self.stacks[i].append(val)
        self.right = max(self.right, i)

        # A partially filled stack remains the next insertion candidate.
        if len(self.stacks[i]) == self.capacity:
            heapq.heappop(self.heap)

    def pop(self) -> int:
        return self.popAtStack(self.right)

    def popAtStack(self, index: int) -> int:
        if not 0 <= index < len(self.stacks):
            return -1

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

        # Only full stacks are absent from the heap, preventing duplicates.
        if len(stack) == self.capacity:
            heapq.heappush(self.heap, index)

        val = stack.pop()

        # Another scan across an empty position requires pushes to refill it.
        while self.right >= 0 and not self.stacks[self.right]:
            self.right -= 1

        return val

Let \(s\) be the number of allocated stacks and \(p\) the number of currently stored plates.

Time: \(O(\log(s + 1))\) amortized per operation; a removal can take \(O(s)\) for the pointer scan.
Space: \(O(p + s)\).

Comment on my solution¶

Your min-heap idea is correct; the index-to-heap-position map is unnecessary because no arbitrary heap deletion is needed.
Popping the candidate index and reinserting it when still nonfull works, but keeping it at the heap root avoids unnecessary heap operations.
The draft is incomplete and currently invalid: helper methods lack self, self.[j] is invalid syntax, and len(self.stacks.append) should be len(self.stacks).
Your pop() must skip empty trailing stacks and register full-to-nonfull transitions in the heap; retaining allocated stacks simplifies this bookkeeping.

# Ideas:
# - difficulty is to find in better than O(n) the left most stack with
#   free space
#   - What could be an invariant for this?
#     - maybe a min heap of stack indices
#       - but what happens when we popAtStack(i) (or pop())
#         - we must update/add i in min heap
#           - so we need also index from "stack idx" -> heap idx
#   - maybe we don't need this index from "stack idx" -> heap idx
#     - because we're always working with the head of the heap

class DinnerPlates:
    def __init__(self, capacity: int):
        self.capacity = capacity
        self.stacks = []
        self.heap = []
        self.index = {} # stack idx -> heap idx

    def _swap(i, j):
        self.heap[i], self.[j] = self.heap[j], self.heap[i]
        self.index[heap[i]] = i
        self.index[heap[j]] = j

    def _sift_up(i):
        while i:
            j = (i - 1) // 2
            if self.heap[i] >= self.heap[j]:
                break
            self._swap(i, j)
            i = j

    def _sift_down(i):
        n = len(self.heap)
        while 2*i + 1 < n:
            j = 2*i + 1
            if j + 1 < n and self.heap[j + 1] < self.heap[j]:
                j += 1
            if self.heap[i] <= self.heap[j]:
                break
            self._swap(i,j)
            i = j

    def _heappop():
        pass
    def _heappush():
        pass

    def push(self, val: int) -> None:
        if self.heap:
            stack_idx = self._heappop()
            del self.index[stack_idx]
        else:
            self.stacks.append([])
            stack_idx = len(self.stacks.append) - 1
        self.stacks[stack_idx].append(val)
        if len(self.stacks[stack_idx]) < self.capacity:
            self._heappush(stack_idx)

    def pop(self) -> int:
        if not self.stacks:
            return -1
        val = self.stacks[-1].pop()
        # TODO: we must update index and heap
        if not self.stacks[-1]:
            self.stacks.pop()
        return val


    def popAtStack(self, index: int) -> int: