🦾

Python: 最小値を返すスタック

2021/08/07に公開

push, pop に加えて min を O(1) でできるスタック
最小値を保持するためのスタックを別の用意しておく(このスタックには一番上に最小値が来るようにする)

  • push するときに、最小値を保持するスタックの一番上の要素と比較し、それよりも小さければ最小値を保持するスタックにも値を加える
  • pop するときに、最小値を保持するスタックの一番上の要素と同じであれば、最小値を保持するスタックも値を削除する
from collections import deque

class StackWithMin:
    def __init__(self):
        self.st = deque([])
        self.st2 = deque([])

    def push(self, value):
        self.st.append(value)
        if (value <= self.min()):
            self.st2.append(value)

    def pop(self):
        value = self.st.pop()
        if (value == self.min()):
            self.st2.pop()
        return value

    def min(self):
        if not self.st2:
            return 1000000000000
        else:
            return self.st2[-1]

test = StackWithMin()
test.push(5)
print(f"stack: {test.st} min: {test.st2[-1]}") # stack: deque([5]) min: 5
test.push(6)
print(f"stack: {test.st} min: {test.st2[-1]}") # stack: deque([5, 6]) min: 5
test.push(3)
print(f"stack: {test.st} min: {test.st2[-1]}") # stack: deque([5, 6, 3]) min: 3
test.push(7)
print(f"stack: {test.st} min: {test.st2[-1]}") # stack: deque([5, 6, 3, 7]) min: 3
test.pop()
print(f"stack: {test.st} min: {test.st2[-1]}") # stack: deque([5, 6, 3]) min: 3
test.pop()
print(f"stack: {test.st} min: {test.st2[-1]}") # stack: deque([5, 6]) min: 5

Discussion