05 用两个栈实现队列

题目描述

用两个栈来实现一个队列，完成队列的Push和Pop操作。队列中的元素为int类型。

解题思路

核心思想：栈是"先进后出"，队列是"先进先出"。用两个栈可以模拟队列：

• 栈1用于入队操作
• 栈2用于出队操作

关键步骤：

• Push：直接压入栈1
• Pop：
  1. 如果栈2不为空，直接从栈2弹出
  2. 如果栈2为空，将栈1中所有元素依次弹出并压入栈2
  3. 从栈2弹出元素

时间复杂度：

• Push：O(1)
• Pop：平均O(1)，最坏O(n)（当栈2为空时需要转移所有元素）

空间复杂度：O(n)

Python 实现

class QueueWithTwoStacks:
    def __init__(self):
        self.stack1 = []  # 用于入队
        self.stack2 = []  # 用于出队

    def push(self, val):
        """
        入队操作
        """
        self.stack1.append(val)

    def pop(self):
        """
        出队操作
        """
        if not self.stack2:
            # 如果栈2为空，将栈1中所有元素转移到栈2
            while self.stack1:
                self.stack2.append(self.stack1.pop())

        if not self.stack2:
            raise Exception("Queue is empty")

        return self.stack2.pop()

    def is_empty(self):
        """判断队列是否为空"""
        return not self.stack1 and not self.stack2

    def peek(self):
        """查看队首元素"""
        if not self.stack2:
            while self.stack1:
                self.stack2.append(self.stack1.pop())

        if not self.stack2:
            raise Exception("Queue is empty")

        return self.stack2[-1]


# 测试代码
if __name__ == "__main__":
    queue = QueueWithTwoStacks()

    queue.push(1)
    queue.push(2)
    queue.push(3)

    print(queue.pop())  # 1
    print(queue.pop())  # 2
    print(queue.pop())  # 3

    # 测试交替操作
    queue.push(4)
    print(queue.pop())  # 4
    queue.push(5)
    print(queue.pop())  # 5

图解

初始状态:
stack1: []
stack2: []

Push 1:
stack1: [1]
stack2: []

Push 2:
stack1: [1, 2]
stack2: []

Push 3:
stack1: [1, 2, 3]
stack2: []

Pop (stack2为空，转移元素):
stack1: []
stack2: [3, 2, 1]  # 转移后顺序翻转
弹出: 1

Pop (stack2不为空):
stack1: []
stack2: [3, 2]
弹出: 2