Leetcode 716. Max Stack

Problem

Design a specialized stack data structure that supports standard stack operations plus the ability to efficiently retrieve and remove the maximum element. Your implementation must provide the following operations:

• push(x): Add an integer x to the top of the stack
• pop(): Remove and return the element at the top of the stack
• top(): Return the element at the top of the stack without removing it
• peekMax(): Return the maximum element in the stack without removing it
• popMax(): Remove and return the maximum element from the stack (if there are multiple instances of the maximum value, remove the most recently added one)

The key challenge is handling popMax() efficiently when the maximum element is not at the top of the stack. After removing the maximum, the relative order of remaining elements must be preserved.

Requirements

1. Implement all five operations: push, pop, top, peekMax, and popMax
2. When multiple elements have the maximum value, popMax must remove the most recently pushed one
3. Maintain the stack property (LIFO order) for regular pop operations
4. Handle edge cases such as empty stack operations appropriately

Constraints

• -10^7 ≤ x ≤ 10^7
• At most 10^4 calls will be made to each function
• Calls to pop, top, peekMax, and popMax will only be made when the stack is non-empty
• Optimize for practical performance across all operations

Examples

Example 1:

` Operations: push(5), push(1), push(5) Stack state: [5, 1, 5] (bottom to top)

top() → 5 (most recent element) popMax() → 5 (removes top 5, most recent max) top() → 1 (now top element) peekMax() → 5 (first 5 is still max) pop() → 1 (removes current top) top() → 5 (only element remaining) `

Example 2:

` Operations: push(3), push(7), push(2), push(7) Stack state: [3, 7, 2, 7] (bottom to top)

peekMax() → 7 (maximum value) popMax() → 7 (removes most recent 7) Stack state: [3, 7, 2] top() → 2 (top unchanged) popMax() → 7 (removes remaining 7) Stack state: [3, 2] `

Example 3:

` Operations: push(1), push(2), push(3), push(4) Stack state: [1, 2, 3, 4] (ascending order)

popMax() → 4 popMax() → 3 popMax() → 2 top() → 1 (only element left) `

Hint 1: Data Structure Choice Consider maintaining two synchronized structures: one for stack order and another for tracking maximum values. A regular stack can handle push/pop/top, but you need an efficient way to find and remove the maximum element regardless of its position.

Hint 2: Handling PopMax When you remove the maximum element (which might not be at the top), you need to temporarily store elements above it, remove the max, then restore those elements in their original order. Think about what auxiliary storage structure would make this efficient.

Hint 3: Tracking Maximums To quickly identify the maximum, you could maintain a separate structure sorted by value. However, when there are duplicates, you need to distinguish between them — perhaps by associating each value with a timestamp or sequence number to track insertion order.

Full Solution `` Approach Explanation:

This solution uses a dual data structure approach with lazy deletion:

1. Main Stack: Stores elements as tuples (value, unique_id) to maintain insertion order and distinguish duplicates.
2. Max Heap: A priority queue storing (-value, -counter) tuples. We negate both values because Python's heapq is a min heap, and negating gives us max heap behavior. Negating the counter ensures that among equal values, the most recent one has the smallest negated counter (highest priority).
3. Lazy Deletion: Instead of immediately removing elements from both structures when one operation deletes them, we mark them in a deleted set. Before returning any value, we clean up the top of the relevant structure by removing already-deleted items.
4. Counter: A monotonically increasing value that serves as a unique identifier for each push operation, allowing us to distinguish between duplicate values and determine which was added most recently.

Time Complexity:

• push: O(log n) — heap insertion
• pop: O(1) amortized — might need to clean deleted items, but each element cleaned once
• top: O(1) amortized — same reasoning as pop
• peekMax: O(1) amortized — might clean heap top, but each element cleaned once
• popMax: O(log n) amortized — heap extraction plus potential cleanup

Space Complexity: O(n) where n is the number of elements currently in the stack, including the stack, heap, and deleted set.

Alternative Approach: Another valid solution uses a TreeMap/Sorted Dictionary to maintain values in sorted order along with their positions, but the implementation complexity is higher and performance characteristics are similar.

import heapq

class MaxStack:
    def __init__(self):
        # Main stack to maintain insertion order
        self.stack = []
        # Max heap to track maximum values (use negative for max heap in Python)
        self.max_heap = []
        # Counter to handle duplicates and track insertion order
        self.counter = 0
        # Set to track deleted elements
        self.deleted = set()
    
    def push(self, x: int) -> None:
        # Add element to stack with a unique ID
        self.stack.append((x, self.counter))
        # Add to max heap (negate value for max heap behavior, negate counter for most recent)
        heapq.heappush(self.max_heap, (-x, -self.counter))
        self.counter += 1
    
    def pop(self) -> int:
        # Clean up any already-deleted elements from top of stack
        while self.stack and self.stack[-1][1] in self.deleted:
            self.stack.pop()
        
        # Get top element and mark as deleted
        value, id_num = self.stack.pop()
        self.deleted.add(id_num)
        return value
    
    def top(self) -> int:
        # Clean up any already-deleted elements from top of stack
        while self.stack and self.stack[-1][1] in self.deleted:
            self.stack.pop()
        
        return self.stack[-1][0]
    
    def peekMax(self) -> int:
        # Clean up any already-deleted elements from top of heap
        while self.max_heap and -self.max_heap[0][1] in self.deleted:
            heapq.heappop(self.max_heap)
        
        return -self.max_heap[0][0]
    
    def popMax(self) -> int:
        # Clean up any already-deleted elements from top of heap
        while self.max_heap and -self.max_heap[0][1] in self.deleted:
            heapq.heappop(self.max_heap)
        
        # Get max element and mark as deleted
        neg_value, neg_id = heapq.heappop(self.max_heap)
        value = -neg_value
        id_num = -neg_id
        self.deleted.add(id_num)
        return value