餐盘栈

题目

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

题解

方法一：

思路

若所有非空栈中最大的下标为n-1，变长数组st维护非空的前n个栈。
用平衡树no_full维护最小未满的栈的下标。

对于插入操作

no_full中的最小值便是st中最小的非空栈，注意对于no_full中最小值若大于等于st的长度，则说明no_full中所有值都越界了，可以清空。
然后便是对最小未满栈插入值，若栈满，则移除no_full的最小值。

对于弹出操作

对于下标不在st的范围内或者对应的栈为空，则返回-1。
否则直接删除对应的栈即可，如果删之前是栈满的则需要在no_full中插入对应下标。此外需要将后缀中空栈全部删除。

能获取$O(log(n))$最小值的数据结构还可以用小根堆。

代码

class DinnerPlates {
public:
    vector<vector<int>> st;
    set<int> no_full;
    int c;
    DinnerPlates(int capacity):c(capacity) {

    }
    
    void push(int val) {
        while (no_full.size() && *no_full.begin() >= st.size()) 
            no_full.erase(no_full.begin());
        if (no_full.empty()) { // [ 0, st.size() ) 都满了
            st.emplace_back(vector<int>(1, val));
            if (c > 1) 
                no_full.insert(st.size()-1);
        } else {
            int idx = *no_full.begin();
            st[idx].push_back(val);
            if (st[idx].size() == c) no_full.erase(no_full.begin());
        }
    }
    
    int pop() {
        return popAtStack((int) st.size() - 1);
    }
    
    int popAtStack(int index) {
        if (index < 0 || index >= st.size() || st[index].empty()) return -1;
        int rt = st[index].back();
        st[index].pop_back();
        no_full.insert(index);
        while (st.size() && st.back().empty()) st.pop_back();
        return rt;
    }
};

/**
 * 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);
 */

小根堆

class DinnerPlates {
public:
    vector<vector<int>> st;
    // set<int> no_full;
    priority_queue<int, vector<int>, greater<int>> no_full;
    int c;
    DinnerPlates(int capacity):c(capacity) {

    }
    
    void push(int val) {
        // while (no_full.size() && *no_full.begin() >= st.size()) 
            // no_full.erase(no_full.begin());
        while (no_full.size() && no_full.top() >= st.size())
            no_full.pop();
        if (no_full.empty()) { // [ 0, st.size() ) 都满了
            st.emplace_back(vector<int>(1, val));
            if (c > 1) 
                // no_full.insert(st.size()-1);
                no_full.push(st.size()-1);
        } else {
            // int idx = *no_full.begin();
            int idx = no_full.top();
            st[idx].push_back(val);
            if (st[idx].size() == c) 
                // no_full.erase(no_full.begin());
                no_full.pop();
        }
    }
    
    int pop() {
        return popAtStack((int) st.size() - 1);
    }
    
    int popAtStack(int index) {
        if (index < 0 || index >= st.size() || st[index].empty()) return -1;
        int rt = st[index].back();
        // no_full.insert(index);
        if (st[index].size() == c) no_full.push(index);
        st[index].pop_back();
        while (st.size() && st.back().empty()) st.pop_back();
        return rt;
    }
};

/**
 * 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);
 */