Taro Logo

Max Stack

Hard
Amazon logo
Amazon
2 views
Topics:
Stacks

Max Stack

Design a data structure that acts like a stack, but also supports retrieving and removing the maximum element in O(1) time (for retrieval) and O(n) for removal in the worst case.

Specifically, implement a MaxStack class with the following methods:

  1. MaxStack(): Initializes an empty max stack.
  2. void push(int x): Pushes element x onto the stack.
  3. int pop(): Removes and returns the element at the top of the stack.
  4. int top(): Returns the element at the top of the stack without removing it.
  5. int peekMax(): Returns the maximum element currently in the stack without removing it. This should be an O(1) operation.
  6. int popMax(): Removes and returns the maximum element currently in the stack. If there are multiple occurrences of the maximum value, only the topmost one should be removed. This operation can take O(n) time in the worst case.

Example Usage:

MaxStack stack = new MaxStack();
stack.push(5);
stack.push(1);
stack.push(5);

System.out.println(stack.top());      // Output: 5
System.out.println(stack.popMax());   // Output: 5
System.out.println(stack.top());      // Output: 1
System.out.println(stack.peekMax());   // Output: 5
System.out.println(stack.pop());      // Output: 1
System.out.println(stack.top());      // Output: 5

Explain the time and space complexity of each operation. Discuss the edge cases and how your implementation handles them.

Solution


Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

  1. What is the expected data type for the elements being stored in the stack, and what is the range of possible values?
  2. Are there any memory constraints or size limitations on the stack's capacity?
  3. What should the `getMax` function return if the stack is empty?
  4. Should the `getMax` function return the most recently added maximum value if multiple elements share the same maximum value?
  5. Is thread safety a concern for the `MaxStack` implementation?

Brute Force Solution

Approach

The brute force method for the Max Stack problem involves simulating every single operation. We create a new stack for each operation performed and store the relevant values in the stack by traversing its contents.

Here's how the algorithm would work step-by-step:

  1. When asked to push a number, create a new stack that is a copy of the current stack, then add the number to the top of the new stack. Make this new stack the current stack.
  2. When asked to pop a number, create a new stack that is a copy of the current stack, then remove the number at the top of the new stack. Make this new stack the current stack. If the stack is empty return null/None/appropriate value.
  3. When asked to get the top number, look at the top of the current stack. If the stack is empty return null/None/appropriate value.
  4. When asked to get the maximum number, inspect every number in the current stack and remember the largest one. If the stack is empty return null/None/appropriate value.
  5. When asked to pop the maximum number, create a new stack that is a copy of the current stack, then inspect every number in the current stack and remember the largest one. Remove the first instance of that largest value from the new stack. Make this new stack the current stack. If the stack is empty return null/None/appropriate value.

Code Implementation

class MaxStack:

    def __init__(self):
        self.stack = []

    def push(self, number):
        # Create a new stack and copy the current stack into it.
        new_stack = self.stack[:]
        new_stack.append(number)
        self.stack = new_stack

    def pop(self):
        # Create a new stack and copy the current stack into it.
        if not self.stack:
            return None
        new_stack = self.stack[:]
        popped_value = new_stack.pop()
        self.stack = new_stack
        return popped_value

    def top(self):
        if not self.stack:
            return None
        return self.stack[-1]

    def peekMax(self):
        # Need to check every number to find the max.
        if not self.stack:
            return None
        maximum_value = self.stack[0]
        for number in self.stack:
            if number > maximum_value:
                maximum_value = number
        return maximum_value

    def popMax(self):
        if not self.stack:
            return None
        new_stack = self.stack[:]
        maximum_value = self.stack[0]
        for number in self.stack:
            if number > maximum_value:
                maximum_value = number

        # Remove the first instance of the maximum value.
        for index, number in enumerate(new_stack):
            if number == maximum_value:
                del new_stack[index]
                break

        self.stack = new_stack
        return maximum_value

Big(O) Analysis

Time Complexity
O(n²)Push and pop operations take O(n) time each because they copy the entire stack. get_max and pop_max operations involve iterating through the entire stack of n elements to find the maximum, also taking O(n) time. In the worst-case scenario, we might perform a sequence of n operations, each potentially taking O(n) time. Therefore, the overall time complexity becomes O(n * n) or O(n²).
Space Complexity
O(N)The brute force method creates a new stack that is a copy of the current stack for each push, pop, and popMax operation. In the worst-case scenario, a series of push operations will result in a stack containing N elements, where N is the number of push operations. Therefore, the space required to store the stack grows linearly with the number of elements, leading to O(N) auxiliary space complexity.

Optimal Solution

Approach

To efficiently manage a stack and quickly find the maximum value within it, we use an auxiliary stack to track the maximums seen so far. This avoids re-calculating the maximum every time a value is requested and provides a constant time lookup.

Here's how the algorithm would work step-by-step:

  1. Keep two stacks: a main stack for storing all numbers, and an auxiliary stack specifically for tracking maximums.
  2. When a number is added to the main stack, check if it's greater than or equal to the current top of the maximums stack. If it is, also add it to the maximums stack. If the maximums stack is empty, add the new number.
  3. When a number is removed from the main stack, check if it's equal to the top of the maximums stack. If it is, also remove it from the maximums stack.
  4. To find the maximum number, simply look at the top of the maximums stack. This gives you the maximum in constant time.

Code Implementation

class MaxStack:

    def __init__(self):
        self.main_stack = []
        self.maximums_stack = []

    def push(self, number):
        self.main_stack.append(number)

        # Only push onto maximums stack if larger than current max.
        if not self.maximums_stack or number >= self.maximums_stack[-1]:

            self.maximums_stack.append(number)

    def pop(self):
        if not self.main_stack:
            return None

        popped_value = self.main_stack.pop()

        # Ensure maximums stack is consistent after popping.
        if self.maximums_stack and popped_value == self.maximums_stack[-1]:

            self.maximums_stack.pop()

        return popped_value

    def top(self):
        if not self.main_stack:
            return None
        return self.main_stack[-1]

    def peekMax(self):
        # Maximums stack top is always the current maximum.
        if self.maximums_stack:

            return self.maximums_stack[-1]
        return None

    def popMax(self):
        maximum_value = self.peekMax()
        if maximum_value is None:
            return None

        temp_stack = []
        while self.top() != maximum_value:
            temp_stack.append(self.pop())

        self.pop()

        # Repopulate main stack with values from temp stack.
        while temp_stack:

            self.push(temp_stack.pop())

        return maximum_value

Big(O) Analysis

Time Complexity
O(1)The push, pop, and getMax operations all involve a constant number of operations. The push operation involves comparing the element to be added with the top of the maximums stack. The pop operation involves comparing the element to be removed with the top of the maximums stack. The getMax operation simply returns the top of the maximums stack. Therefore, each operation takes constant time, regardless of the number of elements (n) in the stack.
Space Complexity
O(N)The algorithm uses an auxiliary stack (maximums stack) that, in the worst case, could store all the elements pushed onto the main stack if they are in increasing order. Therefore, the maximums stack could potentially grow to the same size as the input stack. If N is the number of elements pushed onto the main stack, the auxiliary stack could require space proportional to N. Thus the space complexity is O(N).

Edge Cases

CaseHow to Handle
Stack is initially emptyCheck for empty stack before any pop or peek operations to avoid errors.
Pushing the maximum integer valueEnsure no integer overflow occurs when pushing and updating the max stack.
Popping from the stack when only one element is present, and it is the maximum valueUpdate the max stack appropriately when the last element that is also the maximum is popped.
Large number of push and pop operations exceeding memory limitsConsider using a more memory efficient data structure if memory becomes a concern.
Pushing a value equal to current maximumThe max stack should correctly handle multiple identical maximum values.
Pushing extremely small negative numbersThe max stack should handle negative numbers correctly when finding maximum.
Attempting to find max when the stack contains mixed data types (if applicable in the language)Enforce type consistency by only allowing a single data type (e.g., integers) in the stack.
Popping all elements from the stack, leaving it emptyVerify that the max stack is also empty after all elements are popped from the main stack.