Medium

题目描述

设计一个支持在其元素上进行增量操作的栈。

实现 CustomStack 类:

  • CustomStack(int maxSize) 用 maxSize 初始化对象,maxSize 是栈中最多能容纳的元素数量。
  • void push(int x) 如果栈还没达到 maxSize,就将 x 添加到栈顶。
  • int pop() 弹出栈顶元素,并返回栈顶的值,如果栈为空,返回 -1。
  • void increment(int k, int val) 将栈底的 k 个元素的值都增加 val。如果栈中元素总数小于 k,则将栈中的所有元素都增加 val。

示例 1:

输入:
["CustomStack","push","push","pop","push","push","push","increment","increment","pop","pop","pop","pop"]
[[3],[1],[2],[],[2],[3],[4],[5,100],[2,100],[],[],[],[]]
输出:
[null,null,null,2,null,null,null,null,null,103,202,201,-1]
解释:
CustomStack stk = new CustomStack(3); // 栈为空 []
stk.push(1);                          // 栈变为 [1]
stk.push(2);                          // 栈变为 [1, 2]
stk.pop();                            // 返回 2 --> 返回栈顶值 2,栈变为 [1]
stk.push(2);                          // 栈变为 [1, 2]
stk.push(3);                          // 栈变为 [1, 2, 3]
stk.push(4);                          // 栈仍然是 [1, 2, 3],不能添加其他元素使栈大小超过 3
stk.increment(5, 100);                // 栈变为 [101, 102, 103]
stk.increment(2, 100);                // 栈变为 [201, 202, 103]
stk.pop();                            // 返回 103 --> 返回栈顶值 103,栈变为 [201, 202]
stk.pop();                            // 返回 202 --> 返回栈顶值 202,栈变为 [201]
stk.pop();                            // 返回 201 --> 返回栈顶值 201,栈变为 []
stk.pop();                            // 返回 -1 --> 栈为空,返回 -1

提示:

  • 1 <= maxSize, x, k <= 1000
  • 0 <= val <= 100
  • 每种方法 incrementpush 以及 pop 分别最多调用 1000 次

解题思路

解题思路

本题要求设计一个支持增量操作的栈,关键在于如何高效地实现增量操作。

方法一:暴力解法

最直观的想法是用数组模拟栈,增量操作时直接遍历前 k 个元素逐一增加。这种方法简单但增量操作的时间复杂度为 O(k)。

方法二:延迟更新(推荐)

更优的解法是使用延迟更新的思想。我们维护一个额外的增量数组 inc,其中 inc[i] 表示索引 i 及其之前的所有元素需要增加的值。

核心思想:

  1. pushpop 操作保持 O(1) 复杂度
  2. increment 操作时,只需要在 inc[min(k-1, size-1)] 位置记录增量值
  3. pop 操作时,将当前位置的增量值累加到返回值,并将增量值向下传递给前一个位置

这样避免了增量操作时的遍历,所有操作都能保持较好的时间复杂度。

具体实现:

  • 使用数组 stack 存储栈元素
  • 使用数组 inc 存储每个位置的累积增量
  • pop 时需要处理增量传递逻辑

代码实现

class CustomStack {
private:
    vector<int> stack;
    vector<int> inc;
    int maxSize;
    
public:
    CustomStack(int maxSize) : maxSize(maxSize) {
        stack.reserve(maxSize);
        inc.resize(maxSize, 0);
    }
    
    void push(int x) {
        if (stack.size() < maxSize) {
            stack.push_back(x);
        }
    }
    
    int pop() {
        if (stack.empty()) {
            return -1;
        }
        
        int idx = stack.size() - 1;
        int result = stack[idx] + inc[idx];
        
        if (idx > 0) {
            inc[idx - 1] += inc[idx];
        }
        inc[idx] = 0;
        stack.pop_back();
        
        return result;
    }
    
    void increment(int k, int val) {
        if (!stack.empty()) {
            int idx = min(k - 1, (int)stack.size() - 1);
            inc[idx] += val;
        }
    }
};
class CustomStack:

    def __init__(self, maxSize: int):
        self.stack = []
        self.inc = [0] * maxSize
        self.maxSize = maxSize

    def push(self, x: int) -> None:
        if len(self.stack) < self.maxSize:
            self.stack.append(x)

    def pop(self) -> int:
        if not self.stack:
            return -1
        
        idx = len(self.stack) - 1
        result = self.stack[idx] + self.inc[idx]
        
        if idx > 0:
            self.inc[idx - 1] += self.inc[idx]
        self.inc[idx] = 0
        self.stack.pop()
        
        return result

    def increment(self, k: int, val: int) -> None:
        if self.stack:
            idx = min(k - 1, len(self.stack) - 1)
            self.inc[idx] += val
public class CustomStack {
    private List<int> stack;
    private int[] inc;
    private int maxSize;

    public CustomStack(int maxSize) {
        this.maxSize = maxSize;
        this.stack = new List<int>();
        this.inc = new int[maxSize];
    }
    
    public void Push(int x) {
        if (stack.Count < maxSize) {
            stack.Add(x);
        }
    }
    
    public int Pop() {
        if (stack.Count == 0) {
            return -1;
        }
        
        int idx = stack.Count - 1;
        int result = stack[idx] + inc[idx];
        
        if (idx > 0) {
            inc[idx - 1] += inc[idx];
        }
        inc[idx] = 0;
        stack.RemoveAt(idx);
        
        return result;
    }
    
    public void Increment(int k, int val) {
        if (stack.Count > 0) {
            int idx = Math.Min(k - 1, stack.Count - 1);
            inc[idx] += val;
        }
    }
}
var CustomStack = function(maxSize) {
    this.maxSize = maxSize;
    this.stack = [];
    this.inc = [];
};

CustomStack.prototype.push = function(x) {
    if (this.stack.length < this.maxSize) {
        this.stack.push(x);
        this.inc.push(0);
    }
};

CustomStack.prototype.pop = function() {
    if (this.stack.length === 0) {
        return -1;
    }
    
    let idx = this.stack.length - 1;
    let val = this.stack.pop() + this.inc[idx];
    let incVal = this.inc.pop();
    
    if (idx > 0) {
        this.inc[idx - 1] += incVal;
    }
    
    return val;
};

CustomStack.prototype.increment = function(k, val) {
    let idx = Math.min(k, this.stack.length) - 1;
    if (idx >= 0) {
        this.inc[idx] += val;
    }
};

复杂度分析

操作时间复杂度空间复杂度
pushO(1)O(1)
popO(1)O(1)
incrementO(1)O(1)
总体空间复杂度-O(maxSize)