We need a stack that supports the usual operations (push, pop, top) plus two extra ones: finding the maximum element and removing it. The hard operation is popMax. A regular stack only removes from the top, but popMax removes the maximum, which can sit anywhere in the stack. If several elements share the maximum value, it removes the one closest to the top.
A normal stack gives O(1) push, pop, and top. Supporting peekMax on top of that is straightforward: track the running maximum with a second stack, the same idea used in the Min Stack problem. popMax is the difficult part. The maximum can be deep in the stack, the elements above it must keep their order after removal, and the new maximum can be a different value. The conflict between stack ordering and fast max-tracking is the core of this problem.
The question is whether we can support all five operations efficiently, or whether some operation has to be slow.
-10^7 <= x <= 10^7. Values can be negative and span a wide range, so we handle the full signed integer range rather than assuming positive or small values.10^5 calls. An O(n) per popMax solution with O(1) for the other operations is fast enough to pass, since popMax itself amortizes well over a sequence of calls. The optimal solution drives every operation down to O(log n).Use two stacks. The main stack holds all elements in order. The second stack (the "max stack") tracks the running maximum at each level, the same idea used in the Min Stack problem.
This gives O(1) for push, pop, top, and peekMax. The max stack always has the current maximum at its top, because every time we push a value we also push max(value, currentMax) onto the max stack. Popping both together keeps them aligned.
popMax is harder. The maximum is not always at the top of the main stack, so we pop elements off the top one at a time until we reach the maximum, remove it, then push the popped elements back. The push-back step rebuilds the max stack at the same time, so it stays correct.
The implementation is short and correct. When the stack stays small, the linear popMax cost is acceptable.
stack (the main stack) and maxStack (tracks the running maximum).push(x): Push x onto stack. Push max(x, maxStack.peek()) onto maxStack (or just x if maxStack is empty).pop(): Pop from both stack and maxStack. Return the value from stack.top(): Return stack.peek().peekMax(): Return maxStack.peek().popMax(): Get the maximum from maxStack.peek(). Pop elements into a temp stack until you find the max. Remove it. Push everything back using push() to rebuild the max stack.The bottleneck is popMax, which can pop the entire stack to reach a maximum buried at the bottom. The next approach replaces that linear scan with a heap that finds the maximum in O(log n).
A max-heap finds the maximum in O(log n), and a list maintains stack order. The difficulty is keeping the two structures consistent when popMax removes an element from one of them but not the other.
Lazy deletion resolves this. Instead of removing an element from both structures at once, we mark it deleted and skip it later when it surfaces. Each element carries a unique ID, a counter that increments on every push, so we can identify the exact element to delete even when two elements share the same value.
popMax pops the maximum from the heap, records its ID in a deleted set, and leaves the stack untouched. pop and top first discard any top-of-stack entries whose IDs are in the deleted set, then operate on the first surviving entry. Because each element is added to the deleted set at most once and skipped at most once, the cleanup cost amortizes across operations.
(value, id) pairs), a max-heap of (value, id) pairs, a set of deleted IDs, and a counter for generating unique IDs.push(x): Assign a new unique ID. Push (x, id) onto the stack and the heap. Increment the counter.pop(): Skip any top-of-stack entries whose IDs are in the deleted set. Pop the top entry, add its ID to the deleted set, and return the value.top(): Skip any top-of-stack entries whose IDs are in the deleted set. Return the top value.peekMax(): Skip any top-of-heap entries whose IDs are in the deleted set. Return the top value.popMax(): Skip any top-of-heap entries whose IDs are in the deleted set. Pop the top entry, add its ID to the deleted set, and return the value.The heap with lazy deletion gives good amortized performance, but a single pop or popMax can be slow when it has to skip a long run of deleted entries. The next approach removes each element from both structures immediately, so no entry is ever skipped.
Let's think about what each operation truly needs. Push and pop need stack ordering (know what's on top). PeekMax and popMax need fast access to the maximum. PopMax needs to remove an element from the MIDDLE of the stack without disturbing the order.
No single data structure gives us all of this. But two data structures working together can. A doubly linked list maintains stack order and lets us remove any node in O(1) if we have a pointer to it. A balanced BST (TreeMap/SortedDict) lets us find and access the maximum in O(log n).
The bridge between them: the TreeMap maps each value to a list of linked list node references. When we push a value, we add a node to the tail of the linked list and also record a reference to that node in the TreeMap. When we call popMax, we look up the maximum key in the TreeMap, grab the most recently pushed node with that value, remove it from the linked list in O(1), and remove the reference from the TreeMap. No lazy deletion needed.
push(x): Create a new node, append it to the tail of the linked list. In the TreeMap, add this node reference to the list under key x.pop(): Remove the tail node from the linked list. Remove the last node reference from the TreeMap under the node's value. If that list is empty, remove the key.top(): Return the tail node's value.peekMax(): Return the last key in the TreeMap (the maximum).popMax(): Get the last key from the TreeMap. Get the last node reference from that key's list. Remove that node from the doubly linked list. Remove the reference from the TreeMap.