232 Implement Queue using Stacks

题目

Implement Queue using Stacks

题解

题目要求用两个栈模拟一个队列，一个简单的方法是：每次push时都将元素压入到同一个栈（input_stack）中，当需要pop时则将input_stack里的元素全部弹出并压入到另一个栈（output_stack）里，此时output_stack里的顶部元素就是pop需要移除的元素，将其移除后再把output_stack剩余的元素全部倒回input_stack即可。

上面这个方法虽然简单，但是每次pop时都需要将所有元素都倒腾一遍，其时间复杂度为\(O(n)\)，这里我们可以利用output_stack的一个特点进行优化：output_stack里的元素自顶向下就已经按照队列pop的顺序排列好了。因此我们可以去掉把output_stack剩余元素全部倒回input_stack的操作，在pop的时直接从output_stack顶部进行弹出，当output_stack为空时则先把input_stack里的元素全部倒入output_stack。

在进行了优化之后，每个元素都只会有常数次（三次）移动：

1. 压入input_stack
2. 从input_stack转移到output_stack
3. 从output_stack弹出

因此pop操作的均摊时间复杂度从原来的\(O(n)\)优化到了\(O(1)\)。

代码

class MyQueue:

    def __init__(self):
        self.input_stack = []
        self.output_stack = []

    def push(self, x: int) -> None:
        self.input_stack.append(x)

    def _shift(self):
        while len(self.input_stack) != 0:
            self.output_stack.append(self.input_stack.pop())

    def pop(self) -> int:
        if len(self.output_stack) == 0:
            self._shift()
        return self.output_stack.pop()

    def peek(self) -> int:
        if len(self.output_stack) == 0:
            self._shift()
        return self.output_stack[-1]

    def empty(self) -> bool:
        return len(self.input_stack) == 0 and len(self.output_stack) == 0