Implement Queue using Stacks

How it works

Instead of moving elements back after each pop, we can keep them in the second stack. stack1 handles incoming elements (push), while stack2 holds elements in reversed order for popping. When stack2 is empty and we need to pop, we transfer all elements from stack1 to stack2 at once. Each element is moved at most twice (once to stack2, once when popped), giving amortized O(1) per operation.

Steps

1. For push: Push the element onto s1.
2. For pop: If s2 is empty, transfer all elements from s1 to s2. Then pop from s2.
3. For peek: If s2 is empty, transfer all elements from s1 to s2. Then return the top of s2.
4. For empty: Return true if both s1 and s2 are empty.

Complexity

• Time complexity:
  • O(1) time for initialization.
  • O(1) time for each push() and empty() function calls.
  • O(1) amortized time for each pop() and peek() function calls.
• Space complexity: O(n)

Things to watch for

Checking Only One Stack for Empty

In the optimized two-stack approach, the empty() method must check if both stacks are empty, not just one. A common mistake is returning s1.isEmpty() without considering s2. Elements may exist in s2 ready to be popped even when s1 is empty. The correct check is s1.isEmpty() && s2.isEmpty().

Transferring Elements Back After Every Pop

In the brute force approach, forgetting to transfer elements back from stack2 to stack1 after a pop operation breaks subsequent operations. However, in the optimized approach, the opposite mistake is problematic: unnecessarily transferring elements back defeats the amortized O(1) complexity. The key insight is that elements in s2 are already in the correct order for popping and should remain there until exhausted.

class MyQueue:

    def __init__(self):
        self.s1 = []
        self.s2 = []

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

    def pop(self) -> int:
        if not self.s2:
            while self.s1:
                self.s2.append(self.s1.pop())
        return self.s2.pop()

    def peek(self) -> int:
        if not self.s2:
            while self.s1:
                self.s2.append(self.s1.pop())
        return self.s2[-1]

    def empty(self) -> bool:
        return max(len(self.s1), len(self.s2)) == 0

class MyQueue {
private:
    stack<int> s1, s2;

public:
    MyQueue() {}

    void push(int x) {
        s1.push(x);
    }

    int pop() {
        if (s2.empty()) {
            while (!s1.empty()) {
                s2.push(s1.top());
                s1.pop();
            }
        }
        int res = s2.top();
        s2.pop();
        return res;
    }

    int peek() {
        if (s2.empty()) {
            while (!s1.empty()) {
                s2.push(s1.top());
                s1.pop();
            }
        }
        return s2.top();
    }

    bool empty() {
        return s1.empty() && s2.empty();
    }
};