LeetCode 716: Max Stack

Problem Restatement

We need to design a data structure called MaxStack.

It behaves like a normal stack, but also supports operations related to the maximum value.

The required operations are:

If there is more than one maximum element, popMax() must remove the one closest to the top of the stack. The problem also requires top() in O(1) and the other operations in O(log n) time.

Input and Output

Important constraints:

Examples

Example:

stack = MaxStack()
stack.push(5)
stack.push(1)
stack.push(5)
stack.top()
stack.popMax()
stack.top()
stack.peekMax()
stack.pop()
stack.top()

The returned values are:

[None, None, None, None, 5, 5, 1, 5, 1, 5]

Step by step:

The key detail is the first popMax().

There are two 5 values, but the top-most 5 must be removed.

First Thought: Use One Stack

A simple stack can support:

But peekMax() would require scanning the whole stack.

popMax() would also require scanning the stack, removing the maximum, and restoring the other elements.

That gives O(n) time for maximum-related operations.

Problem With One Stack

The hard part is that we need two different access patterns.

A stack needs access to the most recently added element.

A max stack also needs access to the largest value.

Those are different orderings.

For example:

[5, 1, 4, 3]

The top is 3.

The maximum is 5.

So we need one structure for stack order and another structure for value order.

Key Insight

Use two linked data structures together.

First, use a doubly linked list to maintain stack order.

The tail of the list is the top of the stack.

This gives:

Second, use an ordered map from value to nodes.

Conceptually:

value -> list of nodes with that value

The ordered map lets us find the largest value quickly.

The list of nodes for each value preserves insertion order, so the last node in that list is the top-most occurrence of that value.

For Python, the standard library does not include a balanced ordered map. For LeetCode Python, a common accepted approach uses a heap with lazy deletion. That gives the same idea: one structure tracks stack order, and one structure tracks maximum order.

Algorithm

Use these fields:

self.stack
self.heap
self.removed
self.index

Where:

Why use id?

If two values are equal, popMax() must remove the one closest to the top.

The most recently pushed value has the largest id.

In the heap, we store:

(-value, -id)

Python’s heap is a min heap.

So the largest value has the smallest -value.

Among equal values, the largest id has the smallest -id.

Therefore the heap gives the correct popMax() candidate.

Lazy Deletion

The stack and heap both store the same elements.

But removing from the middle of a heap or stack is expensive.

So when an element is removed from one structure, we mark its id as removed.

Before reading from either structure, we clean its top while the top id is marked as removed.

This is called lazy deletion.

For the stack:

while stack top is removed:
    pop it

For the heap:

while heap top is removed:
    pop it

Operation Details

For push(x):

1. Assign a new id.
2. Push (x, id) onto the stack.
3. Push (-x, -id) onto the heap.

For pop():

1. Clean removed elements from the stack top.
2. Pop the stack top.
3. Mark its id as removed.
4. Return its value.

For top():

1. Clean removed elements from the stack top.
2. Return the stack top value.

For peekMax():

1. Clean removed elements from the heap top.
2. Return -heap[0][0].

For popMax():

1. Clean removed elements from the heap top.
2. Pop the heap top.
3. Convert (-value, -id) back to (value, id).
4. Mark its id as removed.
5. Return the value.

Correctness

Each pushed element receives a unique id, so duplicate values are still distinguishable.

The stack stores all elements in push order. After cleaning lazily removed elements from the top, the last remaining stack entry is exactly the current top of the stack. Therefore top() and pop() return the correct stack-top value.

The heap stores every element as (-value, -id). Because Python returns the smallest heap entry, this ordering gives the largest value first. If several elements have the same value, the entry with the largest id comes first, because -id is smaller. Since larger id means more recently pushed, the heap top is exactly the top-most maximum element.

Lazy deletion keeps the two structures synchronized. If an element is removed by pop(), its id is marked removed, so future heap operations skip it. If an element is removed by popMax(), its id is marked removed, so future stack operations skip it. Cleaning before every read or removal ensures each operation observes only active elements.

Therefore all five operations behave as required.

Complexity

Let n be the number of elements currently stored or waiting to be lazily cleaned.

Each stale element is removed from each structure at most once, so the cleanup cost is amortized.

Implementation

import heapq

class MaxStack:

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

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

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

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

    def pop(self) -> int:
        self._clean_stack()

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

        return value

    def top(self) -> int:
        self._clean_stack()

        return self.stack[-1][0]

    def peekMax(self) -> int:
        self._clean_heap()

        return -self.heap[0][0]

    def popMax(self) -> int:
        self._clean_heap()

        neg_value, neg_idx = heapq.heappop(self.heap)

        value = -neg_value
        idx = -neg_idx

        self.removed.add(idx)

        return value

Code Explanation

The main stack stores values in push order:

self.stack = []

The heap stores values by maximum order:

self.heap = []

Removed ids are tracked here:

self.removed = set()

Each pushed value gets a unique id:

self.index = 0

When pushing, we insert into both structures:

self.stack.append((x, self.index))
heapq.heappush(self.heap, (-x, -self.index))

The stack cleanup removes elements already deleted by popMax():

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

The heap cleanup removes elements already deleted by pop():

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

For pop(), after cleaning, the stack top is valid:

value, idx = self.stack.pop()
self.removed.add(idx)
return value

For peekMax(), after cleaning, the heap top is valid:

return -self.heap[0][0]

For popMax(), the heap returns the maximum value and top-most duplicate:

neg_value, neg_idx = heapq.heappop(self.heap)

Then we mark that id removed so the stack will skip it later.

Example Walkthrough

Use:

stack.push(5)
stack.push(1)
stack.push(5)

The stack contains:

[(5, 0), (1, 1), (5, 2)]

The heap contains:

[(-5, -2), (-1, -1), (-5, 0)]

The heap top is (-5, -2), which means value 5 with id 2.

That is the top-most maximum.

So:

stack.popMax()

returns 5 and marks id 2 as removed.

Now the logical stack is:

[5, 1]

When top() is called, _clean_stack() removes (5, 2) from the physical stack, then returns 1.

Testing

def test_max_stack():
    stack = MaxStack()

    stack.push(5)
    stack.push(1)
    stack.push(5)

    assert stack.top() == 5
    assert stack.popMax() == 5
    assert stack.top() == 1
    assert stack.peekMax() == 5
    assert stack.pop() == 1
    assert stack.top() == 5

    stack = MaxStack()

    stack.push(1)
    stack.push(2)
    stack.push(3)

    assert stack.peekMax() == 3
    assert stack.popMax() == 3
    assert stack.top() == 2

    stack.push(5)
    stack.push(5)
    stack.push(4)

    assert stack.popMax() == 5
    assert stack.top() == 4
    assert stack.popMax() == 5
    assert stack.peekMax() == 4

    stack = MaxStack()

    stack.push(-1)
    stack.push(-3)
    stack.push(-2)

    assert stack.peekMax() == -1
    assert stack.popMax() == -1
    assert stack.top() == -2

    print("all tests passed")

test_max_stack()

Test coverage: