设计一个支持增量操作的栈

标签: 设计 数组

难度: Medium

请你设计一个支持对其元素进行增量操作的栈。

实现自定义栈类 CustomStack

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

示例:

输入:
["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],不能添加其他元素使栈大小变为 4
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

Submission

运行时间: 55 ms

内存: 17.1 MB

class CustomStack:

    def __init__(self, maxSize: int):
        #延迟更新增量的栈
        self.st=[0]*maxSize
        self.add=[0]*maxSize
        self.p=-1
        self.maxSize=maxSize


    def push(self, x: int) -> None:
        if self.p<self.maxSize-1:
            self.p+=1
            self.st[self.p]=x


    def pop(self) -> int:
        if self.p>=0:
            res=self.st[self.p]+self.add[self.p]
            if self.p-1>=0:
                self.add[self.p-1]+=self.add[self.p]
            self.st[self.p]=0
            self.add[self.p]=0
            self.p-=1
            return res
        return -1

    def increment(self, k: int, val: int) -> None:
        cur=min(self.p,k-1)
        if cur>=0:
            self.add[cur]+=val



# Your CustomStack object will be instantiated and called as such:
# obj = CustomStack(maxSize)
# obj.push(x)
# param_2 = obj.pop()
# obj.increment(k,val)

Explain

该题解采用了一个栈结构和一个增量数组来实现CustomStack。两个数组`st`和`add`的长度都是`maxSize`。`st`数组用来存储栈中的元素,而`add`数组用来存储对应位置及其以下所有元素的增量。使用一个指针`p`来表示栈顶位置(初始为-1,表示栈空)。在执行`push`操作时,如果栈未满,就将元素加入到栈顶。在`pop`操作中,除了返回栈顶元素,还需要将当前栈顶的增量传递给下一个栈顶元素。`increment`操作则更新`add`数组,使得最多前k个元素的值都增加`val`,但仅更新栈中实际存在的元素的增量。

时间复杂度: O(1)

空间复杂度: O(maxSize)

class CustomStack:

    def __init__(self, maxSize: int):
        # 初始化栈和增量数组以及栈顶指针
        self.st = [0] * maxSize
        self.add = [0] * maxSize
        self.p = -1
        self.maxSize = maxSize


    def push(self, x: int) -> None:
        # 如果栈未满,则添加元素到栈顶
        if self.p < self.maxSize - 1:
            self.p += 1
            self.st[self.p] = x


    def pop(self) -> int:
        # 如果栈非空,则返回并移除栈顶元素,同时处理增量
        if self.p >= 0:
            res = self.st[self.p] + self.add[self.p]
            if self.p - 1 >= 0:
                self.add[self.p - 1] += self.add[self.p]
            self.st[self.p] = 0
            self.add[self.p] = 0
            self.p -= 1
            return res
        return -1 # 空栈返回-1

    def increment(self, k: int, val: int) -> None:
        # 对栈底的最多k个元素增加val
        cur = min(self.p, k - 1)
        if cur >= 0:
            self.add[cur] += val # 只更新最深影响的增量位置

Explore

在构造函数中初始化`add`数组的目的是为了记录每个栈元素及其以下所有元素的增量信息。这种设计允许在`increment`操作中高效地增加多个栈元素的值,而无需遍历整个栈。在`pop`操作中,为确保增量正确传递,当前栈顶元素的增量(如果有)会被加到返回的栈顶元素上,并且此增量会传递给下一个栈顶元素,即`add[p-1] += add[p]`,然后将当前栈顶的增量清零,这样可以保证增量信息在栈变动时仍然正确和有效。

在`pop`操作中,当前栈顶元素的增量值会被加到下一个栈顶元素的增量中,即`add[p-1] += add[p]`,然后当前栈顶元素的增量被清零。这种设计通常能很好地工作,但可能会在极端情况下导致错误的增量传递,比如当`pop`操作后直接进行`increment`操作且前者未完全清理增量信息时。为避免这种情况,代码确保在每次`pop`时彻底清除当前栈顶的增量。

具体实现方式是,`increment`操作只更新`add`数组中最深受影响的位置的增量值,即`add[min(p, k-1)] += val`。这样的设计是为了优化性能,减少不必要的数组操作。因为在`pop`过程中增量会逐级向上累加,只需在最底层位置更新增量,就可以保证所有相关的栈元素在弹出时都能接收到正确的增量。这种方法减少了操作的复杂性并提高了效率。

当执行连续的`pop`操作时,每次`pop`都会处理增量传递,即将当前栈顶的增量加到下一个栈顶的增量中并清零当前栈顶的增量。这个过程保证了即使连续弹出多个元素,每个元素在弹出时都会加上正确的增量。因此,无论进行多少次`pop`操作,每个元素都会根据`add`数组的记录得到正确的增量更新。