스택 시간복잡도 O(1)으로 최소값 구하기

|

문제출처

문제

How would you design a stack which, in addition to push and pop, also has a function min which returns the minimum element? Push, pop, and min should all operate in O(1) time.

풀이코드

  • O(1) 의 시간복잡도를 가진다는 것은 hash, dictionary 처럼 변수에 직접 접근한다는 의미
  • 변수 mins와 min을 추가하여 최소값에 바로 접근 할 수 있도록 한다.
class Stack():
    def __init__(self):
        self.items = []
        self.mins = []
        self.min = None

    def push(self, item):
        self.items.append(item)
        self.mins.append(self.min)
        if self.min is None or self.min > item:
            self.min = item

    def pop(self):
        self.items.pop()
        self.min = self.mins.pop()

    def get_min(self):
        return self.min

    def peak(self):
        return self.items[-1]

if __name__ == '__main__':
    stack = Stack()
    stack.push(7)
    stack.push(4)
    stack.push(5)
    stack.push(3)
    print(stack.get_min()) # 3
    stack.pop()
    print(stack.peak()) # 5
    stack.pop()
    print(stack.peak()) # 4
    print(stack.get_min()) # 4