Solutions

1. Two Balanced Trees

from sortedcontainers import SortedList

class MaxStack:

    def __init__(self):
        self.stack = SortedList()
        self.values = SortedList()
        self.cnt = 0

    def push(self, x: int) -> None:
        self.stack.add((self.cnt, x))
        self.values.add((x, self.cnt))
        self.cnt += 1

    def pop(self) -> int:
        idx, val = self.stack.pop()
        self.values.remove((val, idx))
        return val

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

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

    def popMax(self) -> int:
        val, idx = self.values.pop()
        self.stack.remove((idx, val))
        return val

Time & Space Complexity

Time Complexity: $O(\log N)$ for each operation except for initialization. All operations other than initialization involve finding/inserting/removing elements in a balanced tree once or twice. In general, the upper bound of time complexity for each of them is $O(\log N)$ . However, note that top and peekMax operations, which require only the last element in a balanced tree, can be done in $O(1)$ with set::rbegin() in C++ and special handling on the last element of SortedList in Python. However, last for TreeSet in Java hasn't implemented similar optimization yet, so we have to get the last element in $O(\log N)$ .
Space complexity: $O(N)$ the maximum size of the two balanced trees.

Where $N$ is the number of elements to add to the stack.

2. Heap + Lazy Update

class MaxStack:

    def __init__(self):
        self.heap = []
        self.cnt = 0
        self.stack = []
        self.removed = set()

    def push(self, x: int) -> None:
        heapq.heappush(self.heap, (-x, -self.cnt))
        self.stack.append((x, self.cnt))
        self.cnt += 1

    def pop(self) -> int:
        while self.stack and self.stack[-1][1] in self.removed:
            self.stack.pop()

        num, idx = self.stack.pop()
        self.removed.add(idx)

        return num

    def top(self) -> int:
        while self.stack and self.stack[-1][1] in self.removed:
            self.stack.pop()

        return self.stack[-1][0]

    def peekMax(self) -> int:
        while self.heap and -self.heap[0][1] in self.removed:
            heapq.heappop(self.heap)

        return -self.heap[0][0]

    def popMax(self) -> int:
        while self.heap and -self.heap[0][1] in self.removed:
            heapq.heappop(self.heap)
            
        num, idx = heapq.heappop(self.heap)
        self.removed.add(-idx)

        return -num

Time & Space Complexity

Time Complexity:
- push: $O(\log N)$ . It costs $O(\log N)$ to add an element to the heap and $O(1)$ to add it to the stack.
- The amortized time complexity of operations caused by a single pop / popMax call is $O(\log N)$ . For a pop call, we first remove the last element in the stack and add its ID to removed in $O(1)$ , resulting in the deletion of the top element in the heap in the future (when peekMax or popMax is called), which has a time complexity of $O(\log N)$ . Similarly, popMax needs $O(\log N)$ immediately and $O(1)$ for the operations later. Note that because we lazy-update the two data structures, future operations might never happen in some cases. However, even in the worst cases, the upper bound of the amortized time complexity is still only $O(\log N)$ .
- top: $O(1)$ , excluding the time cost related to popMax calls we discussed above.
- peekMax: $O(1)$ , excluding the time cost related to pop calls we discussed above.
Space Complexity: $O(N)$ , the maximum size of the heap, stack, and removed.

Where $N$ is the number of elements to add to the stack.