由两个栈组成的队列

题目描述

用两个栈实现队列，支持队列的基本操作。

输入描述

第一行输入一个整数N，表示对队列进行的操作总数。

下面N行每行输入一个字符串S，表示操作的种类。

如果S为"add"，则后面还有一个整数X表示向队列尾部加入整数X。

如果S为"poll"，则表示弹出队列头部操作。

如果S为"peek"，则表示询问当前队列中头部元素是多少。

输出描述

对于每一个为"peek"的操作，输出一行表示当前队列中头部元素是多少。

思路

使用两个栈，一个用于push，一个用于pop，但当pop栈不为空的时候，pop栈的栈顶就是队列头部，可以直接操作；

当pop栈为空的时候，我们需要一次性把push栈中的元素压入pop中。

代码

class Queue_based_stack(object):
    def __init__(self):
        self.push_stack = []
        self.pop_stack = []
        
    def add(self,num):
        self.push_stack.append(num)
        
    def poll(self):
        if self.pop_stack:
            return self.pop_stack.pop()
        elif self.push_stack:
            while (self.push_stack):
                self.pop_stack.append(self.push_stack.pop())
            return self.pop_stack.pop()
        else:
            return None
        
    def peek(self):
        if self.pop_stack:
            return self.pop_stack[-1]
        elif self.push_stack:
            while (self.push_stack):
                self.pop_stack.append(self.push_stack.pop())
            return self.pop_stack[-1]
        else:
            return None
        
queues = Queue_based_stack()
N = int(input())
for i in range(N):
    flag = input().split()
    if flag[0] == 'add':
        queues.add(flag[1])
    elif flag[0] == 'poll':
        queues.poll()
    else:
        print(queues.peek())