get_max() 함수를 갖는 커스텀 Stack 클래스 만들기

• stack은 push, pop, top 함수를 갖는다.
• push는 자료구조의 끝에 요소를 추가, pop은 자료구조의 끝에서 요소를 제거, top은 자료구조 끝의 요소를 조회하는 함수로, 각각 O(1) 시간 안에 수행된다.
• 하지만 stack 구조에서 특정값을 찾으려면 O(n) 시간이 걸린다.
• stack의 최댓값을 리턴하는 get_max() 함수를 작성하고, 시간복잡도를 O(n)보다 낮게 하는 방식을 생각해 보자. => max값을 위한 변수를 하나 더 만들어야 한다.
• 단, 같은 값이 여러번 들어오는 경우도 생각하여 max_stack에 count 변수를 포함한 튜플을 저장하자.
• assert로 유닛테스트를 실행하자.

class CustomStack:
    def __init__(self):
        self._stack = []
        self._max_stack = []  # [(value, count)]

    def push(self, e):
        self._stack.append(e)
        if not self._max_stack:
            self._max_stack.append((e, 1))
        else:
            current_max, count = self._max_stack[-1]
            if e > current_max:
                self._max_stack.append((e, 1))
            elif e == current_max:
                self._max_stack[-1] = (current_max, count + 1)
        # e < current_max인 경우 아무것도 안 함

    def pop(self):
        if not self._stack:
            raise IndexError("stack is empty")
        val = self._stack.pop()
        current_max, count = self._max_stack[-1]
        if val == current_max:
            if count == 1:
                self._max_stack.pop()
            else:
                self._max_stack[-1] = (current_max, count - 1)
        return val

    def top(self):
        if not self._stack:
            raise IndexError("stack is empty")
        return self._stack[-1]

    def get_max(self):
        if not self._stack:
            raise IndexError("stack is empty")
        return self._max_stack[-1][0]
    def __init__(self):
        self.stack = []
        self.max_stack = []
	
    def push(self, v):
        self.stack.append(v)
        if len(self.max_stack) == 0 or v > self.max_stack[-1]:
            self.max_stack.append(v)


    def pop(self):
        if self.stack[-1] in self.max_stack:
            self.max_stack.pop()
        self.stack.pop()
        
        
    def top(self):
        return self.stack[-1]
        
        
    def get_max(self):
        return self.max_stack[-1]


cs = custom_stack()

cs.push(1)
assert cs.stack==[1]

cs.push(2)
assert cs.stack==[1,2]

cs.push(3)
assert cs.stack==[1,2,3]

cs.pop()
assert cs.stack==[1,2]

assert cs.top()==2

assert cs.get_max()==2