LeetCode 1172. 餐盘栈

1172. 餐盘栈

我们把无限数量 ∞ 的栈排成一行，按从左到右的次序从 0 开始编号。每个栈的的最大容量 capacity 都相同。

实现一个叫「餐盘」的类 DinnerPlates：

• DinnerPlates(int capacity) - 给出栈的最大容量 capacity。
• void push(int val) - 将给出的正整数 val 推入 从左往右第一个 没有满的栈。
• int pop() - 返回 从右往左第一个 非空栈顶部的值，并将其从栈中删除；如果所有的栈都是空的，请返回 -1。
• int popAtStack(int index) - 返回编号 index 的栈顶部的值，并将其从栈中删除；如果编号 index 的栈是空的，请返回 -1。

示例：

输入： 
["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],[],[],[],[],[]]
输出：
[null,null,null,null,null,null,2,null,null,20,21,5,4,3,1,-1]

解释：
DinnerPlates D = DinnerPlates(2);  // 初始化，栈最大容量 capacity = 2
D.push(1);
D.push(2);
D.push(3);
D.push(4);
D.push(5);         // 栈的现状为：    2  4
                                    1  3  5
                                    ﹈ ﹈ ﹈
D.popAtStack(0);   // 返回 2。栈的现状为：      4
                                          1  3  5
                                          ﹈ ﹈ ﹈
D.push(20);        // 栈的现状为：  20  4
                                   1  3  5
                                   ﹈ ﹈ ﹈
D.push(21);        // 栈的现状为：  20  4 21
                                   1  3  5
                                   ﹈ ﹈ ﹈
D.popAtStack(0);   // 返回 20。栈的现状为：       4 21
                                            1  3  5
                                            ﹈ ﹈ ﹈
D.popAtStack(2);   // 返回 21。栈的现状为：       4
                                            1  3  5
                                            ﹈ ﹈ ﹈ 
D.pop()            // 返回 5。栈的现状为：        4
                                            1  3 
                                            ﹈ ﹈  
D.pop()            // 返回 4。栈的现状为：    1  3 
                                           ﹈ ﹈   
D.pop()            // 返回 3。栈的现状为：    1 
                                           ﹈   
D.pop()            // 返回 1。现在没有栈。
D.pop()            // 返回 -1。仍然没有栈。

提示：

• 1 <= capacity <= 20000
• 1 <= val <= 20000
• 0 <= index <= 100000
• 最多会对 push，pop，和 popAtStack 进行 200000 次调用。

---

提示 1

Use a data structure to save the plate status. You may need to operate the exact index. Maintain the leftmost vacant stack and the rightmost non-empty stack.

---

提示 2

Use a list of stack to store the plate status. Use heap to maintain the leftmost and rightmost valid stack.

解法1：数组 + 有序集合模拟

用一个数组 stack 来模拟栈，编号为 index 的栈的顶部 stackTop 在数组中的下标 pos 可以通过公式来表示：pos=index*capacity+stackTop。用一个有序集合 poppedPos 来保存被方法 popAtStack 删除的位置。用数组 top 记录每个栈的栈顶元素在栈中的位置，比如 top[1]=2 就表示，编号为 1 的栈，栈顶元素在栈中的下标为 2（从 0 开始计数，capacity>2），即在这个栈中，它上面没有元素，它下面还有两个元素。

执行 push 时，先考虑 poppedPos 中的位置，如果非空，则找出最小的位置，把元素 push 到这个位置。如果为空，则往 stack 后追加，然后更新 top。

执行 popAtStack 时，先找出这个栈现在的栈顶位置，然后把这个位置的元素的下标计算出来并更新栈顶位置，把下标放入 poppedPos，返回元素的值。当然有可能这个栈是空的，此时要返回 −1。

执行 pop 时，情况比较复杂。直观的想法是，返回 stack 元素即可。但 stack 末尾元素可能早已被 popAtStack 删除，因此，应该返回处于 stack 末尾且不位于 popAtStack 中的位置，如果 stack 末尾位置出现在 popAtStack 中，直接把 stack 末尾元素删除，再进行重复判断，直到满足上述条件或者 stack 为空。找到符合条件的位置后，需要更新 top，然后返回元素的值。在执行 pop 时，上述判断可能会执行多次，但是一次 popAtStack 最多带来一次判断，因此不会带来时间复杂度的变大。

我们可以使用一个数组来模拟栈，同时使用有序集合来维护被删除的栈的位置。以下是详细的题解步骤：

1. 初始化：首先，我们需要初始化一个类 DinnerPlates，它包含一个最大容量 capacity 和两个数组 stack 和 top。stack 用来存储栈中元素的值，top 用来记录每个栈的栈顶元素在 stack 中的位置。
2. push 操作：在执行 push 操作时，我们需要将元素 val 推入从左到右第一个没有满的栈。如果 poppedPos（被删除的栈的位置集合）不为空，我们从 poppedPos 中取出最小的索引 index，将元素推入到该索引对应的栈中。如果 poppedPos 为空，我们直接将元素追加到 stack 的末尾，并更新对应栈的 top 值。
3. popAtStack 操作：执行 popAtStack 时，我们首先检查索引 index 对应的栈是否为空。如果不为空，我们找到该栈的栈顶元素在 stack 中的位置，将其从 stack 中删除，并将该索引放入 poppedPos 中，然后返回栈顶元素的值。如果栈为空，则返回 -1。
4. pop 操作：执行 pop 操作时，我们需要找到从右往左第一个非空栈的栈顶元素并将其删除。这需要我们从 stack 的末尾开始检查，找到第一个不在 poppedPos 中的位置。如果找到了，我们更新对应栈的 top 值并返回该元素的值。如果 stack 为空，则返回 -1。

Java版：

class DinnerPlates {
    private int capacity;
    private List<Integer> stack;
    private List<Integer> top;
    private TreeSet<Integer> poppedPos;

    public DinnerPlates(int capacity) {
        this.capacity = capacity;
        stack = new LinkedList<>();
        top = new ArrayList<>();
        poppedPos = new TreeSet<>();
    }
    
    public void push(int val) {
        if (poppedPos.isEmpty()) {
            int pos = stack.size();
            stack.add(val);
            if (pos % capacity == 0) {
                top.add(0);
            } else {
                int index = pos / capacity;
                int stackTop = top.get(index);
                top.set(index, stackTop + 1);
            }
        } else {
            int pos = poppedPos.pollFirst();
            stack.set(pos, val);
            int index = pos / capacity;
            int stackTop = top.get(index);
            top.set(index, stackTop + 1);
        }
    }
    
    public int pop() {
        while (!stack.isEmpty() && poppedPos.contains(stack.size() - 1)) {
            int pos = stack.size() - 1;
            poppedPos.remove(pos);
            stack.remove(pos);
        }
        if (stack.isEmpty()) {
            return -1;
        }
        int pos = stack.size() - 1;
        int index = pos / capacity;
        if (pos % capacity == 0) {
            top.remove(index);
        } else {
            int stackTop = top.get(index);
            top.set(index, stackTop - 1);
        }
        return stack.remove(pos);
    }
    
    public int popAtStack(int index) {
        if (index >= top.size()) {
            return -1;
        }
        int stackTop = top.get(index);
        if (stackTop < 0) {
            return -1;
        }
        top.set(index, stackTop - 1);
        int pos = index * capacity + stackTop;
        poppedPos.add(pos);
        return stack.get(pos);
    }
}

/**
 * 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版：

from sortedcontainers import SortedSet
class DinnerPlates:

    def __init__(self, capacity: int):
        self.capacity = capacity
        self.stack = []
        self.top = []
        self.poppedPos = SortedSet()

    def push(self, val: int) -> None:
        if not self.poppedPos:
            pos = len(self.stack)
            self.stack.append(val)
            if pos % self.capacity == 0:
                self.top.append(0)
            else:
                self.top[-1] += 1
        else:
            pos = self.poppedPos.pop(0)
            self.stack[pos] = val 
            index = pos // self.capacity 
            self.top[index] += 1


    def pop(self) -> int:
        while self.stack and self.poppedPos and self.poppedPos[-1] == len(self.stack) - 1:
            self.poppedPos.pop()
            self.stack.pop()
        if not self.stack:
            return -1 
        pos = len(self.stack) - 1
        if pos % self.capacity == 0:
            self.top.pop()
        else:
            self.top[-1] -= 1
        return self.stack.pop()

    def popAtStack(self, index: int) -> int:
        if index >= len(self.top):
            return -1
        stackTop = self.top[index]
        if stackTop < 0:
            return -1 
        self.top[index] -= 1
        pos = index * self.capacity + stackTop
        self.poppedPos.add(pos)
        return self.stack[pos]


# 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)

复杂度分析

• 时间复杂度：push 操作的时间复杂度为 O(1)，因为每次操作都是直接在数组的末尾添加元素或从有序集合中取出最小元素。popAtStack 操作的时间复杂度为 O(logn)，其中 n 是栈的数量，因为可能需要从有序集合中添加或删除元素。pop 操作的均摊时间复杂度也为 O(logn)，因为每次 pop 操作可能需要从 stack 数组的末尾删除元素，并且可能需要更新 top 数组。
• 空间复杂度：O(n)，其中 n 是 push 调用的次数，因为我们需要存储所有元素以及维护被删除栈的位置集合。