Python实现O(1)时间复杂度的栈（含push、pop、increase操作）

需求是实现一个支持以下三个操作的栈，且所有操作时间复杂度必须为O(1)：

• push(x)：将元素x追加到栈末尾
• pop()：弹出并返回栈顶元素
• increase(k)：将栈顶k个元素的值各加1（例如栈为[1,2,3]，k=2时变为[1,3,4]）

最初的错误实现 #

下面是最初编写的代码，它无法满足O(1)时间复杂度要求，且存在运行报错问题：

from collections import deque


class Stack:
    def __init__(self):
        self.s = []
        self.new = []

    def push(self, x):
        self.s.append(x)
        self.new.append(0)

    def pop(self):
        return self.new.pop() + self.s.pop()

    def increase(self, k):
        b = self.new[-k:]
        a = [1] * k
        self.new = self.new[0:-k]
        self.new += map(sum, zip(a,b))

if __name__ == "__main__":
    s = Stack()
    s.push(1)
    s.push(2)
    s.push(3)
    s.increase(2)
    print(s.pop()) # 4
    s.push(1)
    s.increase(2)
    print(s.pop()) # 2
    print(s.pop()) # 4
    print(s.pop()) # 1

这个实现的问题在于increase操作中，切片、数组创建以及map(sum, zip(...))都是O(k)时间复杂度的操作，完全不符合要求；另外当k大于当前栈元素数量时，切片操作会引发索引错误。

最优O(1)时间复杂度实现 #

下面是符合要求的最优实现，核心思路是无需实际修改栈内元素，只需保证pop()返回结果符合操作后的预期：

class Stack:
    def __init__(self):
        self.s = []
        self.new = [0]

    def push(self, x):
        self.s.append(x)
        self.new.append(0)

    def pop(self):
        self.new[-2] += self.new[-1]
        return self.new.pop() + self.s.pop()

    def increase(self, k):
        self.new[-1] += 1
        self.new[-1-k] -= 1

实现原理说明

这里利用了差分数组的思想：

1. self.s存储原始栈元素，self.new作为辅助差分数组，记录增量变化
2. increase(k)操作时，只需在差分数组对应位置做加减标记：给栈顶k个元素加1，等价于在差分数组的栈顶位置加1，同时在第len(self.s)-k个位置（k个元素的前一个位置）减1，这是O(1)操作
3. pop()操作时，将当前增量传递给前一个元素（self.new[-2] += self.new[-1]），然后弹出当前增量并和原始元素相加得到正确返回值，这一步也是O(1)操作

内容的提问来源于stack exchange，提问作者m__