Dinner Plate Stacks

Problem

You have an infinite number of stacks arranged in a row and numbered (left to right) from 0, each of the stacks has the same maximum capacity.

Implement the DinnerPlates class:

• DinnerPlates(int capacity) Initializes the object with the maximum capacity of the stacks capacity.
• void push(int val) Pushes the given integer val into the leftmost stack with a size less than capacity.
• int pop() Returns the value at the top of the rightmost non-empty stack and removes it from that stack, and returns -1 if all the stacks are empty.
• int popAtStack(int index) Returns the value at the top of the stack with the given index index and removes it from that stack or returns -1 if the stack with that given index is empty.

Examples

Example 1

**Input**
["DinnerPlates", "push", "push", "push", "push", "push", "popAtStack", "push", "push", "popAtStack", "popAtStack", "pop", "pop", "pop", "pop", "pop"]
[[2], [1], [2], [3], [4], [5], [0], [20], [21], [0], [2], [], [], [], [], []]
**Output**
[null, null, null, null, null, null, 2, null, null, 20, 21, 5, 4, 3, 1, -1]

Explanation: 
DinnerPlates D = DinnerPlates(2);  // Initialize with capacity = 2
D.push(1);
D.push(2);
D.push(3);
D.push(4);
D.push(5);         // The stacks are now:  2  4
                                           1  3  5
                                           ﹈ ﹈ ﹈
D.popAtStack(0);   // Returns 2.  The stacks are now:     4
                                                       1  3  5
                                                       ﹈ ﹈ ﹈
D.push(20);        // The stacks are now: 20  4
                                           1  3  5
                                           ﹈ ﹈ ﹈
D.push(21);        // The stacks are now: 20  4 21
                                           1  3  5
                                           ﹈ ﹈ ﹈
D.popAtStack(0);   // Returns 20.  The stacks are now:     4 21
                                                        1  3  5
                                                        ﹈ ﹈ ﹈
D.popAtStack(2);   // Returns 21.  The stacks are now:     4
                                                        1  3  5
                                                        ﹈ ﹈ ﹈ 
D.pop()            // Returns 5.  The stacks are now:      4
                                                        1  3 
                                                        ﹈ ﹈  
D.pop()            // Returns 4.  The stacks are now:   1  3 
                                                        ﹈ ﹈   
D.pop()            // Returns 3.  The stacks are now:   1 
                                                        ﹈   
D.pop()            // Returns 1.  There are no stacks.
D.pop()            // Returns -1.  There are still no stacks.

Constraints

• 1 <= capacity <= 2 * 10^4
• 1 <= val <= 2 * 10^4
• 0 <= index <= 10^5
• At most 2 * 10^5 calls will be made to push, pop, and popAtStack.

Solution

Method 1 – Heap and Set for Push/Pop Management

Intuition

We need to efficiently find the leftmost stack with space for push and the rightmost non-empty stack for pop. We use a min-heap to track available stacks for push and a set to track non-empty stacks for pop.

Approach

1. Use a list of stacks, a min-heap for available stacks (push), and a set for non-empty stacks (pop).
2. For push(val), push to the leftmost available stack (from the heap). If no available, create a new stack.
3. For pop(), pop from the rightmost non-empty stack (from the set). Remove empty stacks from the end.
4. For popAtStack(idx), pop from the stack at idx if not empty. If after pop the stack is not full, add idx to the heap.
5. Maintain heap and set to always point to valid stacks.

Code

C++

class DinnerPlates {
    int cap;
    vector<stack<int>> stacks;
    set<int> nonEmpty;
    priority_queue<int, vector<int>, greater<int>> avail;
public:
    DinnerPlates(int capacity) : cap(capacity) {}
    void push(int val) {
        while (!avail.empty() && avail.top() < stacks.size() && stacks[avail.top()].size() == cap) avail.pop();
        int idx = avail.empty() ? stacks.size() : avail.top();
        if (idx == stacks.size()) stacks.emplace_back();
        stacks[idx].push(val);
        nonEmpty.insert(idx);
        if (stacks[idx].size() < cap) avail.push(idx);
    }
    int pop() {
        while (!nonEmpty.empty() && (stacks[*nonEmpty.rbegin()].empty())) nonEmpty.erase(--nonEmpty.end());
        if (nonEmpty.empty()) return -1;
        int idx = *nonEmpty.rbegin();
        int val = stacks[idx].top(); stacks[idx].pop();
        if (stacks[idx].empty()) nonEmpty.erase(idx);
        if (stacks[idx].size() < cap) avail.push(idx);
        return val;
    }
    int popAtStack(int idx) {
        if (idx >= stacks.size() || stacks[idx].empty()) return -1;
        int val = stacks[idx].top(); stacks[idx].pop();
        if (stacks[idx].empty()) nonEmpty.erase(idx);
        if (stacks[idx].size() < cap) avail.push(idx);
        return val;
    }
};

Go

import "container/heap"
type IntHeap []int
func (h IntHeap) Len() int { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *IntHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *IntHeap) Pop() interface{} { old := *h; x := old[len(old)-1]; *h = old[:len(old)-1]; return x }

type DinnerPlates struct {
    cap int
    stacks [][]int
    avail *IntHeap
    nonEmpty map[int]struct{}
}

func Constructor(capacity int) DinnerPlates {
    h := &IntHeap{}
    heap.Init(h)
    return DinnerPlates{cap: capacity, avail: h, nonEmpty: map[int]struct{}{}}
}

func (d *DinnerPlates) Push(val int) {
    for d.avail.Len() > 0 && (*d.avail)[0] < len(d.stacks) && len(d.stacks[(*d.avail)[0]]) == d.cap {
        heap.Pop(d.avail)
    }
    idx := len(d.stacks)
    if d.avail.Len() > 0 {
        idx = (*d.avail)[0]
    }
    if idx == len(d.stacks) {
        d.stacks = append(d.stacks, []int{})
    }
    d.stacks[idx] = append(d.stacks[idx], val)
    d.nonEmpty[idx] = struct{}{}
    if len(d.stacks[idx]) < d.cap {
        heap.Push(d.avail, idx)
    }
}

func (d *DinnerPlates) Pop() int {
    for len(d.nonEmpty) > 0 {
        maxIdx := -1
        for idx := range d.nonEmpty {
            if idx > maxIdx {
                maxIdx = idx
            }
        }
        if maxIdx == -1 || len(d.stacks[maxIdx]) == 0 {
            delete(d.nonEmpty, maxIdx)
            continue
        }
        val := d.stacks[maxIdx][len(d.stacks[maxIdx])-1]
        d.stacks[maxIdx] = d.stacks[maxIdx][:len(d.stacks[maxIdx])-1]
        if len(d.stacks[maxIdx]) == 0 {
            delete(d.nonEmpty, maxIdx)
        }
        if len(d.stacks[maxIdx]) < d.cap {
            heap.Push(d.avail, maxIdx)
        }
        return val
    }
    return -1
}

func (d *DinnerPlates) PopAtStack(idx int) int {
    if idx >= len(d.stacks) || len(d.stacks[idx]) == 0 {
        return -1
    }
    val := d.stacks[idx][len(d.stacks[idx])-1]
    d.stacks[idx] = d.stacks[idx][:len(d.stacks[idx])-1]
    if len(d.stacks[idx]) == 0 {
        delete(d.nonEmpty, idx)
    }
    if len(d.stacks[idx]) < d.cap {
        heap.Push(d.avail, idx)
    }
    return val
}

Java

import java.util.*;
class DinnerPlates {
    int cap;
    List<Deque<Integer>> stacks = new ArrayList<>();
    TreeSet<Integer> nonEmpty = new TreeSet<>();
    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();
        int idx = avail.isEmpty() ? stacks.size() : avail.peek();
        if (idx == stacks.size()) stacks.add(new ArrayDeque<>());
        stacks.get(idx).push(val);
        nonEmpty.add(idx);
        if (stacks.get(idx).size() < cap) avail.offer(idx);
    }
    public int pop() {
        while (!nonEmpty.isEmpty() && stacks.get(nonEmpty.last()).isEmpty()) nonEmpty.pollLast();
        if (nonEmpty.isEmpty()) return -1;
        int idx = nonEmpty.last();
        int val = stacks.get(idx).pop();
        if (stacks.get(idx).isEmpty()) nonEmpty.remove(idx);
        if (stacks.get(idx).size() < cap) avail.offer(idx);
        return val;
    }
    public int popAtStack(int idx) {
        if (idx >= stacks.size() || stacks.get(idx).isEmpty()) return -1;
        int val = stacks.get(idx).pop();
        if (stacks.get(idx).isEmpty()) nonEmpty.remove(idx);
        if (stacks.get(idx).size() < cap) avail.offer(idx);
        return val;
    }
}

Kotlin

import java.util.*
class DinnerPlates(capacity: Int) {
    private val cap = capacity
    private val stacks = mutableListOf<ArrayDeque<Int>>()
    private val nonEmpty = TreeSet<Int>()
    private val avail = PriorityQueue<Int>()
    fun push(v: Int) {
        while (avail.isNotEmpty() && avail.peek() < stacks.size && stacks[avail.peek()].size == cap) avail.poll()
        val idx = if (avail.isEmpty()) stacks.size else avail.peek()
        if (idx == stacks.size) stacks.add(ArrayDeque())
        stacks[idx].push(v)
        nonEmpty.add(idx)
        if (stacks[idx].size < cap) avail.offer(idx)
    }
    fun pop(): Int {
        while (nonEmpty.isNotEmpty() && stacks[nonEmpty.last()].isEmpty()) nonEmpty.pollLast()
        if (nonEmpty.isEmpty()) return -1
        val idx = nonEmpty.last()
        val v = stacks[idx].pop()
        if (stacks[idx].isEmpty()) nonEmpty.remove(idx)
        if (stacks[idx].size < cap) avail.offer(idx)
        return v
    }
    fun popAtStack(idx: Int): Int {
        if (idx >= stacks.size || stacks[idx].isEmpty()) return -1
        val v = stacks[idx].pop()
        if (stacks[idx].isEmpty()) nonEmpty.remove(idx)
        if (stacks[idx].size < cap) avail.offer(idx)
        return v
    }
}

Python

import heapq
class DinnerPlates:
    def __init__(self, capacity: int):
        self.cap = capacity
        self.stacks = []
        self.avail = []
        self.nonempty = set()
    def push(self, val: int) -> None:
        while self.avail and self.avail[0] < len(self.stacks) and len(self.stacks[self.avail[0]]) == self.cap:
            heapq.heappop(self.avail)
        idx = self.avail[0] if self.avail else len(self.stacks)
        if idx == len(self.stacks):
            self.stacks.append([])
        self.stacks[idx].append(val)
        self.nonempty.add(idx)
        if len(self.stacks[idx]) < self.cap:
            heapq.heappush(self.avail, idx)
    def pop(self) -> int:
        while self.nonempty and (not self.stacks[max(self.nonempty)]):
            self.nonempty.remove(max(self.nonempty))
        if not self.nonempty:
            return -1
        idx = max(self.nonempty)
        val = self.stacks[idx].pop()
        if not self.stacks[idx]:
            self.nonempty.remove(idx)
        if len(self.stacks[idx]) < self.cap:
            heapq.heappush(self.avail, idx)
        return val
    def popAtStack(self, idx: int) -> int:
        if idx >= len(self.stacks) or not self.stacks[idx]:
            return -1
        val = self.stacks[idx].pop()
        if not self.stacks[idx]:
            self.nonempty.remove(idx)
        if len(self.stacks[idx]) < self.cap:
            heapq.heappush(self.avail, idx)
        return val

Rust

use std::collections::{BTreeSet, BinaryHeap};
struct DinnerPlates {
    cap: usize,
    stacks: Vec<Vec<i32>>,
    avail: BinaryHeap<std::cmp::Reverse<usize>>,
    nonempty: BTreeSet<usize>,
}
impl DinnerPlates {
    fn new(capacity: i32) -> Self {
        Self { cap: capacity as usize, stacks: vec![], avail: BinaryHeap::new(), nonempty: BTreeSet::new() }
    }
    fn push(&mut self, val: i32) {
        while let Some(&std::cmp::Reverse(idx)) = self.avail.peek() {
            if idx < self.stacks.len() && self.stacks[idx].len() == self.cap {
                self.avail.pop();
            } else {
                break;
            }
        }
        let idx = if let Some(&std::cmp::Reverse(idx)) = self.avail.peek() { idx } else { self.stacks.len() };
        if idx == self.stacks.len() {
            self.stacks.push(vec![]);
        }
        self.stacks[idx].push(val);
        self.nonempty.insert(idx);
        if self.stacks[idx].len() < self.cap {
            self.avail.push(std::cmp::Reverse(idx));
        }
    }
    fn pop(&mut self) -> i32 {
        while let Some(&idx) = self.nonempty.iter().rev().next() {
            if self.stacks[idx].is_empty() {
                self.nonempty.remove(&idx);
            } else {
                let val = self.stacks[idx].pop().unwrap();
                if self.stacks[idx].is_empty() {
                    self.nonempty.remove(&idx);
                }
                if self.stacks[idx].len() < self.cap {
                    self.avail.push(std::cmp::Reverse(idx));
                }
                return val;
            }
        }
        -1
    }
    fn pop_at_stack(&mut self, idx: usize) -> i32 {
        if idx >= self.stacks.len() || self.stacks[idx].is_empty() {
            return -1;
        }
        let val = self.stacks[idx].pop().unwrap();
        if self.stacks[idx].is_empty() {
            self.nonempty.remove(&idx);
        }
        if self.stacks[idx].len() < self.cap {
            self.avail.push(std::cmp::Reverse(idx));
        }
        val
    }
}

TypeScript

class DinnerPlates {
    private cap: number;
    private stacks: number[][] = [];
    private avail: number[] = [];
    private nonempty: Set<number> = new Set();
    constructor(capacity: number) { this.cap = capacity; }
    push(val: number): void {
        while (this.avail.length && this.avail[0] < this.stacks.length && this.stacks[this.avail[0]].length === this.cap) this.avail.shift();
        const idx = this.avail.length ? this.avail[0] : this.stacks.length;
        if (idx === this.stacks.length) this.stacks.push([]);
        this.stacks[idx].push(val);
        this.nonempty.add(idx);
        if (this.stacks[idx].length < this.cap) this.avail.push(idx);
        this.avail.sort((a, b) => a - b);
    }
    pop(): number {
        while (this.nonempty.size && !this.stacks[Math.max(...this.nonempty)]) this.nonempty.delete(Math.max(...this.nonempty));
        if (!this.nonempty.size) return -1;
        const idx = Math.max(...this.nonempty);
        const val = this.stacks[idx].pop()!;
        if (!this.stacks[idx].length) this.nonempty.delete(idx);
        if (this.stacks[idx].length < this.cap) this.avail.push(idx);
        this.avail.sort((a, b) => a - b);
        return val;
    }
    popAtStack(idx: number): number {
        if (idx >= this.stacks.length || !this.stacks[idx].length) return -1;
        const val = this.stacks[idx].pop()!;
        if (!this.stacks[idx].length) this.nonempty.delete(idx);
        if (this.stacks[idx].length < this.cap) this.avail.push(idx);
        this.avail.sort((a, b) => a - b);
        return val;
    }
}

Complexity

• ⏰ Time complexity: O(log n) per operation, due to heap and set operations.
• 🧺 Space complexity: O(n), for storing the stacks and auxiliary data structures.