[알고리즘, Python] Min Stack 구현하기

오늘은 Min Stack을 구현하는 문제를 풀어보자.

 

문제는 아래 leetcode에서 확인할 수 있다.

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

Min Stack - LeetCode

Min Stack이란 push, pop, top, getMin function을 가지는 스택이다.

 

여기서 전제 조건은 모든 function은 시간 복잡도가 O(1)이어야 한다.

---

pop과 top은 기존 스택처럼 구현하면 된다.

 

push를 할 때도 MinStack이 최소값(minVal이라 하겠다)을 알고 있다면 새로운 값과 비교해서 더 작은 값을 minVal로 바꿔주면 된다.

 

그러나 문제는 pop()을 했을 때 남은 스택에서 최소값이 무엇인지를 구하는 방법이다.

 

기존에 최소값을 구하는 여러가지 방법을 생각해봤지만 아무리 생각해도 O(1)로 구현할 수 있는 방법은 없었다.

 

그래서 결국 힌트를 참고했다. (책 Cracking the Coding Interview 6th Edition Q3.2)

 

세번째 힌트에서는 '각 노드가 substack의 최소값을 알 수 있다면' 라고 적혀 있었다. 여기서 substack이란 노드의 아래에 있는 스택을 의미한다.

서브스택을 대략 그림으로 나타내면 위와 같다.

 

이런 방법으로 구현한 코드는 아래와 같다.

class Node:
    def __init__(self, val, minVal):
        self.val = val
        self.minVal = minVal
        

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

    def push(self, val: int) -> None:
        minVal = self.getMin()
        if val < minVal:
            self.stack.append(Node(val, val))
        else:
            self.stack.append(Node(val, minVal))
            

    def pop(self) -> None:
        self.stack.pop()
        
    def top(self) -> int:
        return self.stack[-1].val

    def getMin(self) -> int:
        if len(self.stack) == 0:
            return float("inf")
        return self.stack[-1].minVal

물론 Node 클래스 사용하지 않고 [val, minVal] 이런식으로 구현해도된다. 

---

출처

본문에 포함