問題3.2

3章入りました。
最初popのこと考えなかったので簡単かと思いきや、、、ってやつですね。
まぁ最初なんでウォーミングアップ。

問題

pushとpopに加えて、最小の要素を返す関数minを持つスタックをどのようにデザインしますか？ただしpush、pop、min関数はすべてO(1)の実行時間になるようにしてください。

pystack.py

#encoding: utf-8


class PyStack():
    '''
    pushとpopを行うクラス。
    >>> stack_instance = PyStack()
    >>> stack_instance.push(10)
    >>> stack_instance.push(20)
    >>> stack_instance.push(0)
    >>> stack_instance.pop()
    0
    >>> stack_instance.pop()
    20
    >>> stack_instance.pop()
    10
    >>> stack_instance.pop()
    '''
    def __init__(self):
        self.stack_data = []

    def push(self, num):
        self.stack_data.append(num)

    def pop(self):
        if len(self.stack_data) > 0:
            return self.stack_data.pop()
        else:
            return None


def _test():
    import doctest
    doctest.testmod()


if __name__ == '__main__':
    _test()

これをインポートしたけど、ほとんど変えてるのでインポートしてる意味ないですね。笑。
まぁ普通のスタックだけクラスで分離しておいた方が、多分後の問題で使う気がするので。

3-2.py

#encoding: utf-8
from pystack import PyStack


class PyStackWithMin(PyStack):
    '''
    スタッククラスに新機能を追加。
    スタックの中で最小値を返す。
    push、pop、min関数の実行時間はすべてO(1)。
    >>> stack_instance = PyStackWithMin()
    >>> stack_instance.min()
    >>> stack_instance.push(30)
    >>> stack_instance.min()
    30
    >>> stack_instance.push(20)
    >>> stack_instance.min()
    20
    >>> stack_instance.push(10)
    >>> stack_instance.min()
    10
    >>> stack_instance.push(10)
    >>> stack_instance.min()
    10
    >>> stack_instance.pop()
    10
    >>> stack_instance.min()
    10
    >>> stack_instance.pop()
    10
    >>> stack_instance.min()
    20
    >>> stack_instance.pop()
    20
    >>> stack_instance.min()
    30
    >>> stack_instance.pop()
    30
    >>> stack_instance.min()
    >>> stack_instance.pop()
    '''
    def __init__(self):
        self.stack_data = []
        self.min_stack = []

    def push(self, num):
        # 最小値更新処理。最小値と同じ値をpushする時の処理に注意。
        if len(self.min_stack) == 0 or num <= self.min_stack[-1]:
            self.min_stack.append(num)
        self.stack_data.append(num)

    def pop(self):
        if len(self.stack_data) > 0:
            # 最小値が取り出されたら、最小値更新。
            if self.stack_data[-1] == self.min_stack[-1]:
                self.min_stack.pop()
            return self.stack_data.pop()
        else:
            return None

    def min(self):
        if len(self.min_stack) == 0:
            return None
        else:
            return self.min_stack[-1]


def _test():
    import doctest
    doctest.testmod()


if __name__ == '__main__':
    _test()