716. Max Stack

Reference: LeetCode
Difficulty: Easy

Problem

Design a max stack that supports push, pop, top, peekMax and popMax.

Note:

• -1e7 <= x <= 1e7
• Number of operations won’t exceed 10000.
• The last four operations won’t be called when stack is empty.

Example:

MaxStack stack = new MaxStack();
stack.push(5); 
stack.push(1);
stack.push(5);
stack.top(); -> 5
stack.popMax(); -> 5
stack.top(); -> 1
stack.peekMax(); -> 5
stack.pop(); -> 1
stack.top(); -> 5

Analysis

My first attempt (ListNode): gist Not Accepted

Failed in case: [3 5 1] (popMax 5, it does not know the max is then 3).

However, by writing the code I improved writing functions in ListNode.

Two Stacks (as in 155. Min Stack)

This solution is based on the brute-force solution in 155. Min Stack.

class MaxStack {  
  Stack<Integer> stack = new Stack<>();
  Stack<Integer> buffer = new Stack<>();
  private int max = Integer.MIN_VALUE;

  public void push(int x) {
    stack.push(x);
    if (x > max) {
      max = x;
    }
  }

  public int pop() { // assume not empty
    int val = stack.pop();
    if (val == max) {
      // update max (find the new max)
      int newMax = Integer.MIN_VALUE; // consider size = 1, min should be reset
      // put aside
      while (stack.size() > 0) {
        int x = stack.pop();
        newMax = Math.max(newMax, x);
        buffer.push(x);
      }
      max = newMax;
      // put it back
      while (buffer.size() > 0) {
        stack.push(buffer.pop());
      }
    }
    return val;
  }

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

  public int peekMax() { return max; }

  public int popMax() {
    int newMax = Integer.MIN_VALUE;
    // find the max
    while (stack.peek() != max) {
      newMax = Math.max(newMax, stack.peek());
      buffer.push(stack.pop());
    }
    // remove max
    stack.pop();
    // still put all to the buffer (we want to find newMax)
    while (stack.size() > 0) {
      newMax = Math.max(newMax, stack.peek());
      buffer.push(stack.pop());
    }
    // put back
    while (buffer.size() > 0) {
      stack.push(buffer.pop());
    }
    int oldMax = max;
    max = newMax;
    return oldMax;
  } 
}

Time: pop() and popMax() take $O(N)$
Space: $O(N)$

Two Stacks

class MaxStack {

  private Stack<Integer> stack = new Stack<>();
  private Stack<Integer> maxStack = new Stack<>();
  private Stack<Integer> buffer = new Stack<>();

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

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

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

  public int popMax() {
    int max = maxStack.peek();
    while (stack.peek() != max) {
      buffer.push(pop()); // pop from two stacks
    }
    pop(); // remove max
    while (buffer.size() > 0) {
      push(buffer.pop()); // push into two stacks
    }
    return max;
  }
}

Note:

The following popMax() is not correct. Why? Use stack.pop() or pop()? Or use which push()?

// bug!
public int popMax() {
  while (stack.peek() != maxStack.peek()) {
    buffer.push(stack.pop());
  }
  stack.pop(); // remove max
  while (buffer.size() > 0) {
    stack.push(buffer.pop());
  }
  return maxStack.pop();
}

Time: Only popMax() takes $O(N)$
Space: $O(N)$

The optimization can be made in finding the largest element by TreeMap.

Nice Intuition: