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

news2024/11/26 23:39:40

优先队列

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

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

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

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

对于 $p o p$ 操作，它等价于 $p o p A tSt a c k$ 操作，将参数 $in d e x$ 设为数组最后位置。
对于 $p o p A tSt a c k$ 操作，简单模拟即可。

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

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;
        }
    }
};