Problem Description

The "Max Stack" problem asks you to design a stack-like data structure that supports the following operations efficiently:

• push(x): Push element x onto the stack.
• pop(): Remove the element on top of the stack and return it.
• top(): Get the element on the top of the stack without removing it.
• peekMax(): Retrieve the maximum element in the stack without removing it.
• popMax(): Retrieve and remove the maximum element in the stack. If there is more than one maximum element, only remove the top-most one (i.e., the one closest to the top of the stack).

The goal is to implement all these operations as efficiently as possible. You may assume all values are integers and that the stack is never empty when pop, top, peekMax, or popMax are called.

Thought Process

At first glance, this problem seems like a straightforward extension of a regular stack. However, the peekMax and popMax operations complicate things because a standard stack cannot efficiently find or remove the maximum element.

A brute-force approach would be to search through the stack for the maximum value each time peekMax or popMax is called, but this would be slow (O(n) time per operation).

To optimize, we need a way to:

• Track the current maximum value efficiently.
• Remove the top-most maximum element efficiently.

Solution Approach

To efficiently support all operations, we combine a regular stack with a second data structure to track maximums:

• Stack: We use a basic stack (e.g., a list or linked list) to store the elements in order.
• Max Stack: We maintain a second stack that keeps track of the maximum element at each level. When we push a new value, we also push the maximum of the new value and the previous maximum onto the max stack.

Here's how each operation works:

1. push(x): Push x onto the main stack. Also push max(x, current max) onto the max stack.
2. pop(): Pop from both the main stack and the max stack.
3. top(): Return the top of the main stack.
4. peekMax(): Return the top of the max stack.
5. popMax(): To remove the top-most maximum, we need to pop elements from the main stack until we find the maximum. We can use a buffer stack to temporarily hold popped elements. Once the maximum is removed, we push the buffered items back onto the main stack, updating the max stack accordingly.

This approach ensures that push, pop, top, and peekMax are all O(1), while popMax is O(n) in the worst case.

Example Walkthrough

Let's walk through a sequence of operations:

• push(5): Main stack: [5], Max stack: [5]
• push(1): Main stack: [5, 1], Max stack: [5, 5]
• push(5): Main stack: [5, 1, 5], Max stack: [5, 5, 5]
• top(): Returns 5
• popMax(): Max is 5. Pop elements until we find 5:
  • Pop 5 (main stack: [5, 1], max stack: [5, 5])
  Return 5.
• top(): Returns 1
• peekMax(): Returns 5
• pop(): Pops 1 (main stack: [5], max stack: [5])
• top(): Returns 5

This shows how the max stack tracks the current maximum and how popMax() removes only the top-most maximum.

Time and Space Complexity

Brute-force approach:

• Each peekMax or popMax requires scanning the entire stack: O(n) per operation.
• Other operations are O(1).

• push, pop, top, peekMax: O(1)
• popMax: O(n) in the worst case (if the max is at the bottom).
• Space: O(n) for both main and max stacks, as each element is stored once in each.

This is a significant improvement, especially for frequent peekMax and pop operations.

Summary

The Max Stack problem is a classic example of using auxiliary data structures to enhance the functionality of a basic stack. By maintaining a second stack to track maximums, we can efficiently support constant-time maximum queries and standard stack operations. The only non-O(1) operation is popMax, which still performs well for most use cases. This approach elegantly balances simplicity and efficiency.

Code Implementation

class MaxStack:
    def __init__(self):
        self.stack = []
        self.maxStack = []

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

    def pop(self) -> int:
        self.maxStack.pop()
        return self.stack.pop()

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

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

    def popMax(self) -> int:
        max_val = self.peekMax()
        buffer = []
        while self.top() != max_val:
            buffer.append(self.pop())
        self.pop()  # pop the max
        for x in reversed(buffer):
            self.push(x)
        return max_val
      

class MaxStack {
    stack<int> st;
    stack<int> maxSt;
public:
    MaxStack() {}

    void push(int x) {
        st.push(x);
        if (maxSt.empty()) maxSt.push(x);
        else maxSt.push(max(x, maxSt.top()));
    }

    int pop() {
        int val = st.top();
        st.pop();
        maxSt.pop();
        return val;
    }

    int top() {
        return st.top();
    }

    int peekMax() {
        return maxSt.top();
    }

    int popMax() {
        int maxVal = peekMax();
        stack<int> buffer;
        while (top() != maxVal) {
            buffer.push(pop());
        }
        pop(); // remove max
        while (!buffer.empty()) {
            push(buffer.top());
            buffer.pop();
        }
        return maxVal;
    }
};
      

import java.util.*;

class MaxStack {
    Stack<Integer> stack;
    Stack<Integer> maxStack;

    public MaxStack() {
        stack = new Stack<>();
        maxStack = new Stack<>();
    }

    public void push(int x) {
        stack.push(x);
        if (maxStack.isEmpty()) {
            maxStack.push(x);
        } else {
            maxStack.push(Math.max(x, maxStack.peek()));
        }
    }

    public int pop() {
        maxStack.pop();
        return stack.pop();
    }

    public int top() {
        return stack.peek();
    }

    public int peekMax() {
        return maxStack.peek();
    }

    public int popMax() {
        int maxVal = peekMax();
        Stack<Integer> buffer = new Stack<>();
        while (top() != maxVal) {
            buffer.push(pop());
        }
        pop(); // remove max
        while (!buffer.isEmpty()) {
            push(buffer.pop());
        }
        return maxVal;
    }
}
      

var MaxStack = function() {
    this.stack = [];
    this.maxStack = [];
};

MaxStack.prototype.push = function(x) {
    this.stack.push(x);
    if (this.maxStack.length === 0) {
        this.maxStack.push(x);
    } else {
        this.maxStack.push(Math.max(x, this.maxStack[this.maxStack.length - 1]));
    }
};

MaxStack.prototype.pop = function() {
    this.maxStack.pop();
    return this.stack.pop();
};

MaxStack.prototype.top = function() {
    return this.stack[this.stack.length - 1];
};

MaxStack.prototype.peekMax = function() {
    return this.maxStack[this.maxStack.length - 1];
};

MaxStack.prototype.popMax = function() {
    const max = this.peekMax();
    const buffer = [];
    while (this.top() !== max) {
        buffer.push(this.pop());
    }
    this.pop(); // remove max
    while (buffer.length) {
        this.push(buffer.pop());
    }
    return max;
};