This notebook was prepared by [Donne Martin](http://donnemartin.com). Source and license info is on [GitHub](https://github.com/donnemartin/interactive-coding-challenges).

Solution Notebook

Problem: Implement a stack with push, pop, and min methods running O(1) time.

Constraints
Test Cases
Algorithm
Code
Unit Test

Constraints

Can we assume this is a stack of ints?
- Yes
If we call this function on an empty stack, can we return sys.maxsize?
- Yes
Can we assume we already have a stack class that can be used for this problem?
- Yes

Test Cases

Push/pop on empty stack
Push/pop on non-empty stack

Algorithm

We'll use a second stack to keep track of the minimum values.

Min

If the second stack is empty, return an error code (max int value)
Else, return the top of the stack, without popping it

Complexity:

Time: O(1)
Space: O(1)

Push

Push the data
If the data is less than min
- Push data to second stack

Complexity:

Time: O(1)
Space: O(n)

Pop

Pop the data
If the data is equal to min
- Pop the top of the second stack
Return the data

Complexity:

Time: O(1)
Space: O(1)

Code



In [1]:

    
%run ../stack/stack.py



In [2]:

    
import sys


class MyStack(Stack):

    def __init__(self, top=None):
        self.min_vals = Stack()
        super(MyStack, self).__init__(top)

    def min(self):
        if self.min_vals.top is None:
            return sys.maxsize
        else:
            return self.min_vals.peek()

    def push(self, data):
        super(MyStack, self).push(data)
        if data < self.min():
            self.min_vals.push(data)

    def pop(self):
        data = super(MyStack, self).pop()
        if data == self.min():
            self.min_vals.pop()
        return data

Unit Test



In [3]:

    
%%writefile test_stack_min.py
from nose.tools import assert_equal


class TestStackMin(object):

    def test_stack_min(self):
        print('Test: Push on empty stack, non-empty stack')
        stack = MyStack()
        stack.push(5)
        assert_equal(stack.peek(), 5)
        assert_equal(stack.min(), 5)
        stack.push(1)
        assert_equal(stack.peek(), 1)
        assert_equal(stack.min(), 1)
        stack.push(3)
        assert_equal(stack.peek(), 3)
        assert_equal(stack.min(), 1)
        stack.push(0)
        assert_equal(stack.peek(), 0)
        assert_equal(stack.min(), 0)

        print('Test: Pop on non-empty stack')
        assert_equal(stack.pop(), 0)
        assert_equal(stack.min(), 1)
        assert_equal(stack.pop(), 3)
        assert_equal(stack.min(), 1)
        assert_equal(stack.pop(), 1)
        assert_equal(stack.min(), 5)
        assert_equal(stack.pop(), 5)
        assert_equal(stack.min(), sys.maxsize)

        print('Test: Pop empty stack')
        assert_equal(stack.pop(), None)

        print('Success: test_stack_min')


def main():
    test = TestStackMin()
    test.test_stack_min()


if __name__ == '__main__':
    main()









    



Overwriting test_stack_min.py



In [4]:

    
run -i test_stack_min.py









    



Test: Push on empty stack, non-empty stack
Test: Pop on non-empty stack
Test: Pop empty stack
Success: test_stack_min