Implement Queue Using Stacks
Statement
Design a queue using only two stacks. Implement the methods push
, pop
, peek
, and empty
.
Constraints
- 1 ≤
x
≤ 100
Solution
"""
production algorithm
"""
from data_structures.stack.stack import Stack
class CustomQueue:
def __init__(self):
self.stack1 = Stack()
self.stack2 = Stack()
def push(self, x):
"""
time complexity O(n)
space complexity O(1)
"""
if self.stack1.empty():
self.stack1.push(x)
return
while not self.stack1.empty():
self.stack2.push(self.stack1.pop())
self.stack2.push(x)
while not self.stack2.empty():
self.stack1.push(self.stack2.pop())
def pop(self):
"""
time complexity O(1)
space complexity O(1)
"""
return self.stack1.pop()
def peek(self):
"""
time complexity O(1)
space complexity O(1)
"""
return self.stack1.top()
def empty(self):
"""
time complexity O(1)
space complexity O(1)
"""
return self.stack1.empty()
"""
unit tests
"""
from unittest import TestCase
from algorithms.stacks.implement_queue_using_stacks.custom_queue import CustomQueue
class CustomQueueTestCase(TestCase):
def test_push(self):
# Given
queue = CustomQueue()
numbers = [1, 3, 5, 7, 9]
# When
[queue.push(number) for number in numbers]
actual = queue.stack1.list
# Then
expected = [9, 7, 5, 3, 1]
self.assertEqual(actual, expected)
def test_pop(self):
# Given
queue = CustomQueue()
numbers = [1, 3, 5, 7, 9]
[queue.push(number) for number in numbers]
# When
actual = [queue.pop() for _ in range(len(numbers))]
# Then
expected = [1, 3, 5, 7, 9]
self.assertEqual(actual, expected)
def test_peek(self):
# Given
queue = CustomQueue()
numbers = [1, 3, 5, 7, 9]
[queue.push(number) for number in numbers]
# When
actual = queue.peek()
# Then
expected = 1
self.assertEqual(actual, expected)
def test_empty(self):
# Given
queue = CustomQueue()
# When
actual = queue.empty()
# Then
expected = True
self.assertEqual(actual, expected)
Strategy
Stacks.
Explanation
Queue operations pop
, peek
, and empty
rely on one stack. However, in order to keep elements of the queue sorted, a second stack is used. When calling push
, all elements of the first stack are pushed to the second stack and then the new element is pushed to the top of the second stack. Then, all elements of the second stack are pushed back to the first stack. Afterwards, the new element is at the bottom of the first stack.
Time Complexity
The time complexity of the function push
is O(n)
since all elements of the queue are visited. The time complexity of all other functions pop
, peek
, and empty
is O(1)
.
Space Complexity
The auxiliary space complexity of the function push
is O(n)
since all elements of the queue are pushed to a second stack. The auxiliary space complexity of all other functions pop
, peek
, and empty
is O(1)
.