import unittest class Stack(object): def __init__(self): """Initialize an empty stack""" self.items = [] def push(self, item): """Push a new item onto the stack""" self.items.append(item) def pop(self): """Remove and return the last item""" # If the stack is empty, return None # (it would also be reasonable to throw an exception) if not self.items: return None return self.items.pop() def peek(self): """Return the last item without removing it""" if not self.items: return None return self.items[-1] class MaxStack(object): # Implement the push, pop, and get_max methods def __init__(self): self.m = Stack() self.stack = Stack() def push(self, item): if self.m.peek() is None: self.m.push(item) elif item >= self.m.peek(): self.m.push(item) self.stack.push(item) def pop(self): x = self.stack.pop() if x == self.m.peek(): self.m.pop() return x def get_max(self): return self.m.peek() ################################################################################ # Tests ################################################################################ class Test(unittest.TestCase): def test_stack_usage(self): max_stack = MaxStack() max_stack.push(5) actual = max_stack.get_max() expected = 5 self.assertEqual(actual, expected) max_stack.push(4) max_stack.push(7) max_stack.push(7) max_stack.push(8) actual = max_stack.get_max() expected = 8 self.assertEqual(actual, expected) actual = max_stack.pop() expected = 8 self.assertEqual(actual, expected) actual = max_stack.get_max() expected = 7 self.assertEqual(actual, expected) actual = max_stack.pop() expected = 7 self.assertEqual(actual, expected) actual = max_stack.get_max() expected = 7 self.assertEqual(actual, expected) actual = max_stack.pop() expected = 7 self.assertEqual(actual, expected) actual = max_stack.get_max() expected = 5 self.assertEqual(actual, expected) actual = max_stack.pop() expected = 4 self.assertEqual(actual, expected) actual = max_stack.get_max() expected = 5 self.assertEqual(actual, expected) unittest.main(verbosity=2)