716. Max Stack - Detailed Explanation
Problem Statement
Description:
Design a stack that supports the following operations in addition to the usual push and pop:
- 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.
- peekMax() – Retrieve the maximum element in the stack.
- popMax() – Retrieve the maximum element in the stack and remove it. If there is more than one maximum element, only remove the one closest to the top.
Example 1:
- Operations:
push(5)
→ stack becomes[5]
push(1)
→ stack becomes[5, 1]
push(5)
→ stack becomes[5, 1, 5]
top()
→ returns5
popMax()
→ returns5
(removes the top-most 5) and stack becomes[5, 1]
top()
→ returns1
peekMax()
→ returns5
pop()
→ returns1
and stack becomes[5]
Example 2:
- Operations:
push(2)
→ stack:[2]
push(8)
→ stack:[2, 8]
push(3)
→ stack:[2, 8, 3]
peekMax()
→ returns8
popMax()
→ returns8
, stack becomes[2, 3]
top()
→ returns3
Example 3:
- Operations:
push(7)
→ stack:[7]
push(7)
→ stack:[7, 7]
push(3)
→ stack:[7, 7, 3]
popMax()
→ returns7
(removes the top-most 7) and stack becomes[7, 3]
peekMax()
→ returns7
Constraints:
- The number of operations can be large (e.g., up to 10⁴ or 10⁵).
- The elements pushed are integers.
- You must support duplicate values.
- All operations should be as efficient as possible.
Hints
-
Tracking the Maximum:
Think about how you might maintain the current maximum value at each stage of stack operations. Can you use an extra data structure to “remember” the maximum so far? -
Removing a Specific Element:
For thepopMax()
operation, consider how you might locate and remove the maximum element even if it isn’t on the top of the stack. Sometimes a secondary stack (or a doubly linked list with a tree structure) can help by “buffering” elements while you search for the max.
Approach 1: Brute Force Using a Single List
Idea
-
Push, Pop, Top:
These operations work exactly like a regular stack using a list (or array). -
peekMax():
Scan the list to find the maximum element. -
popMax():
Find the maximum value in the list, then use an auxiliary buffer (another list) to pop elements until you find the top-most occurrence of that maximum. Remove it and then push the buffered elements back.
Step-by-Step Walkthrough (Visual Example)
Imagine the following sequence of operations:
- push(5) → Stack:
[5]
- push(1) → Stack:
[5, 1]
- push(5) → Stack:
[5, 1, 5]
- popMax():
- Find maximum:
5
. - Pop elements until you see
5
from the top:- Pop
5
(found max) → Do not push it into the buffer. - Then push back any buffered elements (none in this case).
- Pop
- Final stack becomes
[5, 1]
.
- Find maximum:
Complexity Analysis
- push(x): O(1)
- pop(): O(1)
- top(): O(1)
- peekMax(): O(n) (since you scan through the stack)
- popMax(): O(n) in the worst case (if you have to pop almost all elements)
Python Code (Brute Force)
Java Code (Brute Force)
Approach 2: Optimal Approach Using a Doubly Linked List & TreeMap
Idea
This approach is more advanced and aims to optimize the popMax()
operation by achieving better-than-linear time on average.
-
Doubly Linked List:
Use a doubly linked list to store the stack elements. This structure allows:- O(1) insertion (push) and deletion (pop) at the tail.
- O(1) deletion of a given node when you have a direct pointer to it.
-
TreeMap (Balanced BST):
Maintain a TreeMap (or any balanced BST) where the keys are the values pushed onto the stack and the values are lists of nodes (from the doubly linked list) that have that key. The TreeMap gives you:- O(log n) access to the maximum key.
- O(log n) update time when inserting or deleting nodes.
How It Works
-
push(x):
Create a new node with valuex
and insert it at the end of the doubly linked list. Also, insert a reference to this node in the TreeMap under keyx
. -
pop():
Remove the node at the tail of the doubly linked list and update the TreeMap by removing the corresponding node reference. -
top():
Return the value of the tail node. -
peekMax():
Retrieve the last key from the TreeMap (the maximum value). -
popMax():
Use the TreeMap to find the maximum value. Get the most recently added node (i.e., the one closest to the top) from the list corresponding to that value. Remove that node from both the doubly linked list and the TreeMap.
Complexity Analysis
- push(x): O(1) for the doubly linked list insertion plus O(log n) for updating the TreeMap.
- pop(), top(): O(1)
- peekMax(): O(1) to O(log n) (depending on the TreeMap implementation)
- popMax(): O(log n) (for TreeMap lookup and update) plus O(1) to remove from the linked list.
Python Code (Using Doubly Linked List & SortedDict)
Note: Python’s standard library does not include a built-in balanced BST. One can use the sortedcontainers module’s
SortedDict
for a similar effect. This solution uses that module.
Java Code (Using Doubly Linked List & TreeMap)
Common Mistakes & Edge Cases
-
Forgetting Duplicate Values:
- Many candidates forget that there can be multiple occurrences of the maximum value. Ensure your data structure (or auxiliary buffer) can handle duplicates.
-
Not Updating Both Data Structures:
- When using a two-structure approach (like the doubly linked list with a TreeMap), always update both when performing operations (e.g., remove the node reference from the TreeMap when it’s popped from the list).
-
Edge Cases:
- Empty Stack: Be sure to handle calls to
pop()
,top()
, orpopMax()
when the stack is empty. - Single Element: Ensure your solution works when there’s only one element in the stack.
- Empty Stack: Be sure to handle calls to
Alternative Variations & Related Problems
-
Min Stack:
Similar to Max Stack, design a stack that supports retrieving the minimum element in O(1) time. -
Design a Stack With Increment Operation:
Where you have to implement an additional operation to increment the bottom k elements by a given value. -
Related Problems for Further Practice:
- Sliding Window Maximum
- Valid Parentheses
- Two Sum
- Implement Queue using Stacks
Final Thoughts
When preparing for coding interviews, it’s valuable to start with a simple (brute force) approach to understand the problem, and then move to more optimized data structures for better performance. The two primary solutions discussed above illustrate:
- A simple, easy-to-implement solution using a single list (ideal for understanding the operations).
- A more optimal solution using a doubly linked list combined with a TreeMap (or SortedDict) that efficiently supports the
popMax()
operation.
Understanding the trade-offs (simplicity vs. efficiency) and the design patterns used (auxiliary data structures to track extra information) will help you in many similar coding interview problems.
GET YOUR FREE
Coding Questions Catalog
