1172. Dinner Plate Stacks

You have an infinite number of stacks arranged in a row and numbered (left to right) from 0, each of the stacks has the same maximum capacity.

Implement the DinnerPlates class:

• DinnerPlates(int capacity) Initializes the object with the maximum capacity of the stacks capacity.
• void push(int val) Pushes the given integer val into the leftmost stack with a size less than capacity.
• int pop() Returns the value at the top of the rightmost non-empty stack and removes it from that stack, and returns -1 if all the stacks are empty.
• int popAtStack(int index) Returns the value at the top of the stack with the given index index and removes it from that stack or returns -1 if the stack with that given index is empty.

class DinnerPlates {
    vector<vector<int>> st;
    priority_queue<int, vector<int>, greater<int>> pushAt;
    int cap;
public:
    DinnerPlates(int capacity): cap(capacity) {
        st = vector<vector<int>>();
    }

    void push(int val) {
        while(!pushAt.empty() and st.size() <= pushAt.top()) pushAt.pop();

        if(pushAt.empty()) {
            st.emplace_back();
            st.back().push_back(val);
            if(st.back().size() < cap) pushAt.push(st.size() - 1);
        } else {
            st[pushAt.top()].push_back(val);
            if(st[pushAt.top()].size() == cap) pushAt.pop();
        }
    }

    int pop() {
        while(!st.empty() and st.back().empty()) {
            st.pop_back();
        }
        if(st.empty()) return -1;
        int res = st.back().back();
        if(st.back().size() == cap) pushAt.push(st.size() - 1);
        st.back().pop_back();
        return res;
    }

    int popAtStack(int index) {
        while(!st.empty() and st.back().empty()) {
            st.pop_back();
        }
        if(st.size() <= index) return -1;
        if(st.size() - 1 == index) return pop();
        if(st[index].empty()) return -1;
        int res = st[index].back();
        if(st[index].size() == cap) pushAt.push(index);
        st[index].pop_back();
        return res;
    }
};

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