LeetCode / Min Stack

（ブログ記事からの転載）

[https://leetcode.com/problems/min-stack/]

Design a stack that supports push, pop, top, and retrieving the minimum element in constant time.

push(x) -- Push element x onto stack.
pop() -- Removes the element on top of the stack.
top() -- Get the top element.
getMin() -- Retrieve the minimum element in the stack.

Example:

MinStack minStack = new MinStack();
minStack.push(-2);
minStack.push(0);
minStack.push(-3);
minStack.getMin(); --> Returns -3.
minStack.pop();
minStack.top(); --> Returns 0.
minStack.getMin(); --> Returns -2.

変わり種の問題ですが、push, pop, top, getMinのメソッドを持つクラスを定義する問題です。
クラス・メソッドの定義の仕方が分かっていれば造作もない問題ですが、ちょっとした落とし穴はあります。

解答・解説

解法1

何も考えずに書けば、このようになると思います。

class MinStack(object):

    def __init__(self):
        self.stack = []

    def push(self, x):
        self.stack.append(x)

    def pop(self):
        if len(self.stack) > 0: self.stack.pop()

    def top(self):
        if len(self.stack) == 0: return None
        return self.stack[-1]

    def getMin(self):
        return min(self.stack)

これでもsubmitは通るんですが、getMinメソッドの実行時に$O(n)$の空間計算量が必要になります。しかし実際にはその必要はありません。

解法2

今回は、push, pop, topメソッドで各要素を追加・削除・取得すること以外には、getMinで最小値を取得することだけがやりたい処理です。

そこでpushメソッドでlistの変数stackにappendする値を、(追加した要素, これまでの最小値と追加した要素のうち小さい方の値)のタプルにします。
これにより、getMinではstack[-1][1]で最小値を取得できるので、$O(1)$で処理できることになります。

class MinStack(object):

    def __init__(self):
        self.stack = []

    def push(self, x):
        if len(self.stack) == 0:
            self.stack.append([x,x])
        else:
            self.stack.append([x, min(x, self.stack[-1][1])])

    def pop(self):
        if len(self.stack) > 0: self.stack.pop()

    def top(self):
        if len(self.stack) == 0: return None
        return self.stack[-1][0]

    def getMin(self):
        return self.stack[-1][1]