Design prefix-sum function and max stack

Question

LeetCode 303. Range Sum Query – Immutable (design a reusable prefix-sum function) LeetCode 716. Max Stack (similar design question)

https://leetcode.com/problems/range-sum-query-immutable/description/ https://leetcode.com/problems/max-stack/description/

Implement Max Stack with popMax

Design a function that processes a sequence of operations on a Max Stack. The stack supports the following commands: (1) "push x": push integer x onto the stack; (2) "pop": remove and return the top element; (3) "top": return the top element without removing it; (4) "peekMax": return the current maximum element; (5) "popMax": remove and return the maximum element; if multiple maxima exist, remove the one closest to the top. The function must return a list of integers representing the results of all operations that produce output (pop, top, peekMax, popMax), in order. All inputs are valid; operations that require a non-empty stack will not be called on an empty stack.

Constraints

• 1 <= len(operations) <= 20000
• Operations are one of: "push x", "pop", "top", "peekMax", "popMax"
• -10^9 <= x <= 10^9
• Operations that require a non-empty stack (pop, top, peekMax, popMax) will not be called on an empty stack

Examples

Example 1

Input: operations = ["push 5","push 1","push 5","top","popMax","top","peekMax","pop","top"]
Output: [5,5,1,5,1,5]

Explanation: After pushes: [5,1,5]. top -> 5. popMax removes the topmost 5 -> [5,1]. top -> 1. peekMax -> 5. pop -> 1 -> [5]. top -> 5.

Example 2

Input: operations = ["push 2","push 2","push 1","peekMax","popMax","peekMax","pop","top"]
Output: [2,2,2,1,2]

Explanation: Stack evolves: [2,2,1]. peekMax -> 2. popMax removes the topmost 2 -> [2,1]. peekMax -> 2. pop -> 1 -> [2]. top -> 2.

Hints

1. Maintain a parallel stack that tracks the maximum at each depth.
2. For popMax, move elements to a temporary buffer until the maximum is at the top, remove it, then restore the buffered elements.
3. When restoring elements from the buffer, use the same push logic so the max-tracking structure stays correct.

Approach

Use two stacks: one for values and one for tracking the running maximum at each position. push, pop, top, and peekMax are O(1) by updating/reading the top of both stacks. For popMax, store the current maximum, then pop elements into a temporary buffer until the top equals that maximum; remove it, then push back buffered elements to restore order and recompute running maxima.

• Time complexity: push/top/peekMax/pop: O(1); popMax: O(n) in the worst case
• Space complexity: O(n) for the stacks and temporary buffer