餐盘栈

标签: 栈 设计 哈希表 堆（优先队列）

难度: Hard

我们把无限数量 ∞ 的栈排成一行，按从左到右的次序从 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 次调用。

Submission

运行时间: 438 ms

内存: 91.9 MB

from sortedcontainers import *
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:
                stackPos = len(self.top) - 1
                stackTop = self.top[stackPos]
                self.top[stackPos] = stackTop + 1
        else:
            pos = self.poppedPos.pop(0)
            self.stack[pos] = val
            index = pos // self.capacity
            stackTop = self.top[index]
            self.top[index] = stackTop + 1

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

    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] = stackTop - 1
        pos = index * self.capacity + stackTop
        self.poppedPos.add(pos)
        return self.stack[pos]

Explain

该题解使用了多个栈和一个有序集合来实现餐盘类。主要思路如下： 1. 用一个列表 `stack` 来存储所有的餐盘，每个餐盘用一个整数表示。 2. 用一个列表 `top` 来存储每个栈的栈顶位置。 3. 用一个有序集合 `poppedPos` 来存储被弹出的位置。 4. 对于 `push` 操作，如果有被弹出的位置，就将新餐盘放到最左边的被弹出位置；否则，将新餐盘放到最右边的栈上。 5. 对于 `pop` 操作，从最右边的栈开始弹出餐盘，并更新 `top` 和 `poppedPos`。 6. 对于 `popAtStack` 操作，根据给定的栈编号计算出餐盘在 `stack` 中的位置，将其弹出，并更新 `top` 和 `poppedPos`。

时间复杂度: O(log n)，其中 n 为调用 `push` 的次数。

空间复杂度: O(n)，其中 n 为调用 `push` 的次数。

from sortedcontainers import *
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:
                stackPos = len(self.top) - 1
                stackTop = self.top[stackPos]
                self.top[stackPos] = stackTop + 1
        else:
            # 如果有被弹出的位置，就将新餐盘放到最左边的被弹出位置
            pos = self.poppedPos.pop(0)
            self.stack[pos] = val
            index = pos // self.capacity
            stackTop = self.top[index]
            self.top[index] = stackTop + 1

    def pop(self) -> int:
        while self.stack and self.poppedPos and self.poppedPos[-1] == len(self.stack) - 1:
            # 如果最右边的餐盘已经被弹出，就将其从 stack 和 poppedPos 中移除
            self.stack.pop()
            pos = self.poppedPos.pop()
            if pos % self.capacity == 0:
                self.top.pop()
        if not self.stack:
            return -1
        else:
            pos = len(self.stack) - 1
            val = self.stack[pos]
            self.stack.pop()
            if pos % self.capacity == 0 and self.top:
                self.top.pop()
            elif self.top:
                self.top[-1] -= 1
            return val

    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] = stackTop - 1
        pos = index * self.capacity + stackTop
        self.poppedPos.add(pos)
        return self.stack[pos]

Explore

在 `DinnerPlates` 类的 `push` 操作中，我们使用了一个有序集合 `poppedPos` 来追踪所有被弹出的位置。这些位置代表了栈中那些已经被弹出但尚未被新餐盘占据的空位。当执行 `push` 操作时，首先检查 `poppedPos` 集合是否为空。如果不为空，这意味着有可用的空位，我们从 `poppedPos` 中取出最小的位置（使用 `pop(0)`，这是因为 `SortedSet` 保持元素有序），并将新的餐盘放置在此位置。这样就确保了被弹出的位置在被重新填充前确实是空的。

在 `pop` 操作中，如果最右边的餐盘已经被弹出（即位置在 `poppedPos` 中），那么在弹出最右边的餐盘之前需要先从 `poppedPos` 中移除这些位置。这通过检查 `poppedPos` 的最后一个元素是否等于 `stack` 最后一个元素的索引来完成。如果是，我们从 `stack` 和 `poppedPos` 中同时移除这些元素。此外，如果弹出位置正好是一个栈的开始位置，还需要从 `top` 列表中移除相应的栈顶索引。这个过程确保了 `poppedPos` 中的索引总是有效的，并且 `stack` 中不会留下无效的空洞。

在 `popAtStack` 方法中，我们通过计算指定栈的当前栈顶位置来弹出餐盘。这是通过 `top[index]` 来得知的。在弹出之后，`top[index]` 会减一，表示栈顶向下移动了一个位置。此外，我们使用 `poppedPos` 集合来记录被弹出的位置，确保在随后的操作中这些位置可以被正确管理。即便栈顶的餐盘已被弹出，通过正确维护 `top` 和 `poppedPos`，我们总是能确保返回正确的餐盘值，因为这些值仍然保存在 `stack` 中，只是它们的索引被记录在了 `poppedPos` 中。