design

716. Max Stack.py

class MaxStack:
    """   
    push 1, 5
    tmp_stack = []
    stack = [5,1]
    max_stack = [5,5]
    
    這題多了一個pop max 是最難的地方
    tmp = stakc.pop() # 1
    tmp_max = max_stack.pop() # 5
    當你popmax把5拿掉後, stack和max stack就不再是[5]和[5]需要改成stack=[1], max_stack=[1]
    我
    stack = [5]
    max_stack = [5]
    [5,1,2,3,4]
    [5,5,5,5,5]
    
    tmp_stack = [4,3,2,1]
    再把他tmp_stack的元素塞回去, 維持
    stack = [1,2,3,4]
    max_stack = [1,2,3,4]
    In worst case, for popmax, we need to pop everything from tmp_stack and max_stack. 
    Then push everything back from tmp_stack, which takes O(n)
    T: Excecpt Max stack, all operations take in O(1)
    S: O(n)
    """
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.stack = []
        self.max_stack = [float("-inf")]

    def push(self, x: int) -> None:
        if x >= self.max_stack[-1]:
            self.max_stack.append(x)
        else:
            # 如果x小於top of stack, 我就放最大的元素
            self.max_stack.append(self.max_stack[-1])
        self.stack.append(x)
        
    def pop(self) -> int:
        self.max_stack.pop()
        return self.stack.pop()

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

    def peekMax(self) -> int:
        return self.max_stack[-1]

    def popMax(self) -> int:
        tmp_stack = []
        tmp_max = self.max_stack.pop()
        tmp = self.stack.pop()
        while tmp_max!= tmp:
            tmp_stack.append(tmp)
            tmp = self.stack.pop()
            tmp_max = self.max_stack.pop()
        while tmp_stack:
            self.push(tmp_stack.pop())
        return tmp_max


# Your MaxStack object will be instantiated and called as such:
# obj = MaxStack()
# obj.push(x)
# param_2 = obj.pop()
# param_3 = obj.top()
# param_4 = obj.peekMax()
# param_5 = obj.popMax()