Max Stack

How it works

We can use a standard stack for regular push/pop/top operations and a max heap to quickly find the maximum. The challenge is that removing an element from one structure does not automatically remove it from the other. We solve this with lazy deletion: when we remove an element, we record its index in a removed set. Before accessing the top of the stack or heap, we skip over any elements that have been marked as removed. This defers the actual cleanup until it is needed.

Steps

1. Maintain a stack for LIFO access, a max heap for maximum access, and a removed set for tracking deleted indices.
2. push(x): Add (x, cnt) to both stack and heap. Increment cnt.
3. pop(): Skip elements at the top of stack that are in removed. Pop the top element, add its index to removed, and return the value.
4. top(): Skip removed elements at the top of stack, then return the top value.
5. peekMax(): Skip removed elements at the top of heap, then return the top value.
6. popMax(): Skip removed elements at the top of heap. Pop the top, add its index to removed, and return the value.

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.

Things to watch for

Forgetting to Synchronize Both Data Structures

When using two balanced trees (or a stack + heap), removing an element from one structure without removing the corresponding entry from the other leads to stale data. Both pop and popMax must update both structures to maintain consistency.

Incorrect Handling of Duplicate Values

Multiple elements can have the same value but different insertion indices. Using only the value as a key causes collisions and incorrect behavior. Always pair values with unique identifiers (like insertion count) to distinguish between duplicates.

Not Implementing Lazy Deletion Properly

In the heap + lazy update approach, forgetting to skip removed elements before accessing the top of the stack or heap returns invalid data. The cleanup loop must run before every top, pop, peekMax, and popMax operation.

Using Wrong Comparator for Max Heap

When implementing a max heap, accidentally using a min heap comparator (or vice versa) returns the wrong maximum element. Ensure the comparator orders elements so the largest value and most recent index appear at the top.

Integer Overflow in Counter

If the counter used for unique indices is not managed properly and the stack experiences many push operations, the counter could overflow. While rare in practice, using a sufficiently large integer type prevents this edge case.

import heapq

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

class MaxStack {
private:
    stack<pair<int, int>> stk;
    priority_queue<pair<int, int>> heap;
    unordered_set<int> removed;
    int cnt;

public:
    MaxStack() { cnt = 0; }

    void push(int x) {
        stk.push({x, cnt});
        heap.push({x, cnt});
        cnt++;
    }

    int pop() {
        while (removed.count(stk.top().second)) {
            stk.pop();
        }

        const pair<int, int> p = stk.top();
        stk.pop();
        removed.insert(p.second);

        return p.first;
    }

    int top() {
        while (removed.count(stk.top().second)) {
            stk.pop();
        }

        return stk.top().first;
    }

    int peekMax() {
        while (removed.count(heap.top().second)) {
            heap.pop();
        }

        return heap.top().first;
    }

    int popMax() {
        while (removed.count(heap.top().second)) {
            heap.pop();
        }

        const pair<int, int> p = heap.top();
        heap.pop();
        removed.insert(p.second);

        return p.first;
    }
};