Leetcode 716 Solution

class MaxStack:

    def __init__(self):
        self.values = []
        self.max_values = []

    def push(self, x: int) -> None:
        self.values.append(x)
        if not self.max_values:
            self.max_values.append(x)
        else:
            self.max_values.append(max(x, self.max_values[-1]))

    def pop(self) -> int:
        ans = self.values[-1]
        self.values.pop()
        self.max_values.pop()
        return ans

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

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

    def popMax(self) -> int:
        i = len(self.values) - 1

        elements_to_put_back = []
        while i >= 0 and self.values[i] != self.max_values[i]:
            elements_to_put_back.append((
                self.values.pop(),
                self.max_values.pop(),
            ))
            i -= 1

        ans = self.values[-1]
        self.values.pop()
        self.max_values.pop()

        for value, max_value in reversed(elements_to_put_back):
            self.push(value)

        return ans

# 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()