Leetcode 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 次调用。

解题思路

这个问题首先不难想到暴力法，建立一个list存放我们的栈。

• push：遍历list看那个栈可以插入元素，那么插入即可。如果所有的栈都不能添加元素，那么添加新栈。时间复杂度O(n)。
• pop：从后向前遍历list，看那个栈中的元素可以弹出，那么弹出即可，否则的话返回-1。时间复杂度O(n)。
• popAtStack：直接可以判断index的栈是不是空，不是空最后一个元素出栈，否则返回-1。时间复杂度O(1)。

class DinnerPlates:
    def __init__(self, capacity: int):
        self.data = []
        self.n = capacity

    def push(self, val: int) -> None:
        for i in self.data:
            if i and len(i) < self.n:
                i.append(val)
                break
        else:
            self.data.append([])
            self.data[-1].append(val)

    def pop(self) -> int:
        for i in range(len(self.data)-1, -1, -1):
            if self.data[i]:
                return self.data[i].pop()
        else:
            return -1

    def popAtStack(self, index: int) -> int:
        if self.data[index]:
            return self.data[index].pop()
        return -1

由于最多会对 push，pop，和 popAtStack 进行 200000 次调用，所以O(n)的操作就很慢了。不难想到，我们可以将那些没有装满的栈装进一个数据结构中存储。可以使用的数据结构有堆和红黑树，我这里使用最小堆存储没有装满的栈的位置，因为有序，所以堆顶元素一定是最左边满足条件的位置。

• push：首先判断堆中位置对应的栈是不是满了，我们将满的栈从堆中剔除。如果堆变成空了，说明所有的栈都满了，我们需要添加新栈，并且将新栈加入堆中。最后需要往堆顶的栈中添加新元素即可。
• pop：首先判断栈顶是不是空，如果是空的话，我们需要将之弹出，然后运行popAtStack()。
• popAtStack：首先需要判断输入的index是不是合法，index对应的栈是不是空。如果既合法也不为空的话，我们就可以将index对应的栈弹出，并且需要往堆中添加index（因为此时index对应的栈一定是不满的）。

import heapq
class DinnerPlates:
    def __init__(self, capacity: int):
        self.data = []
        self.q = []
        self.n = capacity

    def push(self, val: int) -> None:
        while self.q and self.q[0] < len(self.data) and len(self.data[self.q[0]]) == self.n:
            heapq.heappop(self.q)
        if not self.q:
            heapq.heappush(self.q, len(self.data))
            self.data.append([])
        self.data[self.q[0]].append(val)

    def pop(self) -> int:
        while self.data and not self.data[-1]:
            self.data.pop()
        return self.popAtStack(len(self.data) - 1)

    def popAtStack(self, index: int) -> int:
        if 0 <= index < len(self.data) and self.data[index]:
            heapq.heappush(self.q, index)
            return self.data[index].pop()
        return -1

通过使用堆的方式，此时push、pop和popAtStack的时间复杂度就变成了O(logn)。

reference:

https://leetcode.com/problems/dinner-plate-stacks/discuss/366331/C%2B%2BPython-Solution

我将该问题的其他语言版本添加到了我的GitHub Leetcode

如有问题，希望大家指出！！！