力扣(LeetCode)1172. 餐盘栈(C++)

优先队列

解题思路：根据题意模拟。用数组存储无限数量的栈。重在实现 $push$ 和 $pop$ 操作。

1. 对于 $push$ 操作，需要知道当前从左往右第一个空栈的下标。分两类讨论：
   ①所有栈都是满的，那么我们创建一个新栈，新栈存 $push$ 进来的值，再将新栈加入数组尾部。
   ②已存在的某些栈未满，那么从左往右遍历数组，找到第一个未满的栈，第一个未满的栈存 $push$ 进来的值。

这样一来，我们发现每个 $push$ 的操作②的最坏平均时间复杂度 $O(n)$ ，一共 $n$ 个 $push$ ，总体时间复杂度 $O(n^2)$，题目给出 $push$ 操作最多调用 $2\times 10^5$ 次，时间 $O(n^2)$ 必然超时。

操作②既然是遍历，有一个直观的优化方法：用优先队列，小根堆存储未满栈的下标。这样以来，维护未满的栈的时间复杂度就是 $O(logn)$，每次维护完毕，堆顶就是最小下标。总体时间复杂度 $O(nlogn)$，可以接受。

2. 对于 $pop$ 操作，它等价于 $popAtStack$ 操作，将参数 $index$ 设为数组最后位置。

3. 对于 $popAtStack$ 操作，简单模拟即可。

提示：对于优先队列的维护，请思考边界。也可以看代码注释。

class DinnerPlates {
public:
    int capacity;
    vector<stack<int>> S;
    priority_queue<int, vector<int>, greater<int>> pq;  // 维护未满栈的下标
    DinnerPlates(int capacity) {
        this->capacity = capacity;
    }
    
    void push(int val) {
        if (pq.size() && pq.top() >= S.size()) {  // 清空越界下标
            pq = priority_queue<int, vector<int>, greater<int>>();
        }
        // while (pq.size() && pq.top() >= S.size()) pq.pop();  // 清空越界下标
        if (pq.empty()) {  // 插入新栈
            stack<int> ts;
            ts.push(val);
            if (capacity > 1) pq.push(S.size());
            S.push_back(ts);
        } else {  // 插入第一个未满栈
            int pos = pq.top();
            S[pos].push(val);
            if (S[pos].size() >= capacity) {  // 插入后，栈满
                pq.pop();  // 下标弹栈
            }
        }
    }
    
    int pop() {
        return popAtStack(S.size() - 1);
    }
    
    int popAtStack(int index) {
        if (index >= S.size() || S[index].empty()) {
            return -1;
        } else {
            int ans = S[index].top();
            S[index].pop();
            if (S[index].size() == capacity - 1) pq.push(index);
            while (S.size() && S.back().empty()) S.pop_back();
            return ans;
        }
    }
};

时间复杂度 $O(nlogn)$ : 一共 $n$ 次操作，每个操作的时间复杂度 $O(logn)$ ，时间瓶颈在于维护堆。总体时间复杂度 $O(nlogn)$
空间复杂度 $O(n)$ : 优先队列、栈的最坏空间复杂度 $O(n)$

AC

致语

• 理解思路很重要
• 读者有问题请留言，清墨看到就会回复的。