Min stack

min_stack.py

class MinStack(object):

    def __init__(self):
        self.main_stack = []
        self.min_stack = []

    def push(self, n):
        if not self.main_stack:
            self.main_stack.append(n)
            self.min_stack.append(n)
        else:
            min_so_far = self.min_stack[-1]
            if n < min_so_far:
                self.main_stack.append(n)
                self.min_stack.append(n)
            else:
                self.main_stack.append(n)
                self.min_stack.append(min_so_far)

    def pop(self):
        if not self.main_stack:
            return

        self.min_stack.pop()
        return self.main_stack.pop()

    def get_min(self):

        if not self.min_stack:
            return
        return self.min_stack[-1]


if __name__ == "__main__":

    stack = MinStack()
    stack.push(5)
    print(stack.get_min())
    stack.push(4)
    print(stack.get_min())
    stack.push(2)
    print(stack.get_min())
    stack.push(3)
    print(stack.get_min())
    stack.pop()
    print(stack.get_min())
    stack.pop()
    print(stack.get_min())