golang 解法, 执行用时: 424 ms, 内存消耗: 64.1 MB, 提交时间: 2023-04-28 08:40:33

type DinnerPlates struct {
    capacity int     // 栈的容量
    stacks   [][]int // 所有栈
    idx      hp      // 最小堆，保存未满栈的下标
}

func Constructor(capacity int) DinnerPlates {
    return DinnerPlates{capacity: capacity}
}

func (d *DinnerPlates) Push(val int) {
    if d.idx.Len() > 0 && d.idx.IntSlice[0] >= len(d.stacks) {
        d.idx = hp{} // 堆中都是越界下标，直接清空
    }
    if d.idx.Len() == 0 { // 所有栈都是满的
        d.stacks = append(d.stacks, []int{val}) // 添加一个新的栈
        if d.capacity > 1 { // 新的栈没有满
            heap.Push(&d.idx, len(d.stacks)-1) // 入堆
        }
    } else { // 还有未满栈
        i := d.idx.IntSlice[0]
        d.stacks[i] = append(d.stacks[i], val) // 入栈
        if len(d.stacks[i]) == d.capacity { // 栈满了
            heap.Pop(&d.idx) // 从堆中去掉
        }
    }
}

func (d *DinnerPlates) Pop() int {
    // 等价为 popAtStack 最后一个非空栈
    return d.PopAtStack(len(d.stacks) - 1)
}

func (d *DinnerPlates) PopAtStack(index int) int {
    if index < 0 || index >= len(d.stacks) || len(d.stacks[index]) == 0 {
        return -1 // 非法操作
    }
    if len(d.stacks[index]) == d.capacity { // 满栈
        heap.Push(&d.idx, index) // 元素出栈后，栈就不满了，把下标入堆
    }
    bk := len(d.stacks[index]) - 1
    val := d.stacks[index][bk]
    d.stacks[index] = d.stacks[index][:bk]
    for len(d.stacks) > 0 && len(d.stacks[len(d.stacks)-1]) == 0 {
        d.stacks = d.stacks[:len(d.stacks)-1] // 去掉末尾的空栈（懒删除，堆中下标在 push 时处理）
    }
    return val
}

type hp struct{ sort.IntSlice }
func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any   { a := h.IntSlice; v := a[len(a)-1]; h.IntSlice = a[:len(a)-1]; return v }


/**
 * Your DinnerPlates object will be instantiated and called as such:
 * obj := Constructor(capacity);
 * obj.Push(val);
 * param_2 := obj.Pop();
 * param_3 := obj.PopAtStack(index);
 */

java 解法, 执行用时: 53 ms, 内存消耗: 147.3 MB, 提交时间: 2023-04-28 08:40:09

class DinnerPlates {
    // 栈的容量
    private int capacity;
    // 所有栈
    private List<Deque<Integer>> stacks = new ArrayList<>();
    // 最小堆，保存未满栈的下标
    private PriorityQueue<Integer> idx = new PriorityQueue<>();

    public DinnerPlates(int capacity) {
        this.capacity = capacity;
    }

    public void push(int val) {
        if (!idx.isEmpty() && idx.peek() >= stacks.size())
            idx.clear(); // 堆中都是越界下标，直接清空
        if (idx.isEmpty()) { // 所有栈都是满的
            var st = new ArrayDeque<Integer>();
            st.push(val);
            stacks.add(st); // 添加一个新的栈
            if (capacity > 1) // 新的栈没有满
                idx.add(stacks.size() - 1); // 入堆
        } else { // 还有未满栈
            var st = stacks.get(idx.peek());
            st.push(val); // 入栈
            if (st.size() == capacity) // 栈满了
                idx.poll(); // 从堆中去掉
        }
    }

    public int pop() {
        // 等价为 popAtStack 最后一个非空栈
        return popAtStack(stacks.size() - 1);
    }

    public int popAtStack(int index) {
        if (index < 0 || index >= stacks.size() || stacks.get(index).isEmpty())
            return -1; // 非法操作
        var st = stacks.get(index);
        if (st.size() == capacity) // 满栈
            idx.add(index); // 元素出栈后，栈就不满了，把下标入堆
        int val = st.pop();
        // 去掉末尾的空栈（懒删除，堆中下标在 push 时处理）
        while (!stacks.isEmpty() && stacks.get(stacks.size() - 1).isEmpty())
            stacks.remove(stacks.size() - 1);
        return val;
    }
}

/**
 * Your DinnerPlates object will be instantiated and called as such:
 * DinnerPlates obj = new DinnerPlates(capacity);
 * obj.push(val);
 * int param_2 = obj.pop();
 * int param_3 = obj.popAtStack(index);
 */

python3 解法, 执行用时: 556 ms, 内存消耗: 94.9 MB, 提交时间: 2023-04-28 08:37:17

class DinnerPlates:
    def __init__(self, capacity: int):
        self.capacity = capacity  # 栈的容量
        self.stacks = []  # 所有栈
        self.h = []  # 最小堆，保存未满栈的下标

    def push(self, val: int) -> None:
        if self.h and self.h[0] >= len(self.stacks):
            self.h = []  # 堆中都是越界下标，直接清空
        if self.h:  # 还有未满栈
            self.stacks[self.h[0]].append(val)  # 入栈
            if len(self.stacks[self.h[0]]) == self.capacity:  # 栈满了
                heappop(self.h)  # 从堆中去掉
        else:  # 所有栈都是满的
            self.stacks.append([val])  # 添加一个新的栈
            if self.capacity > 1:  # 新的栈没有满
                heappush(self.h, len(self.stacks) - 1)  # 入堆

    def pop(self) -> int:
        # 等价为 popAtStack 最后一个非空栈
        return self.popAtStack(len(self.stacks) - 1)

    def popAtStack(self, index: int) -> int:
        if index < 0 or index >= len(self.stacks) or len(self.stacks[index]) == 0:
            return -1  # 非法操作
        if len(self.stacks[index]) == self.capacity:  # 满栈
            heappush(self.h, index)  # 元素出栈后，栈就不满了，把下标入堆
        val = self.stacks[index].pop()
        while self.stacks and len(self.stacks[-1]) == 0:
            self.stacks.pop()  # 去掉末尾的空栈（懒删除，堆中下标在 push 时处理）
        return val

# Your DinnerPlates object will be instantiated and called as such:
# obj = DinnerPlates(capacity)
# obj.push(val)
# param_2 = obj.pop()
# param_3 = obj.popAtStack(index)