diff options
Diffstat (limited to 'data_structures_and_algorithms/largest-stack.py')
-rw-r--r-- | data_structures_and_algorithms/largest-stack.py | 107 |
1 files changed, 107 insertions, 0 deletions
diff --git a/data_structures_and_algorithms/largest-stack.py b/data_structures_and_algorithms/largest-stack.py new file mode 100644 index 000000000000..aab9671eb6d3 --- /dev/null +++ b/data_structures_and_algorithms/largest-stack.py @@ -0,0 +1,107 @@ +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) |