def assert_equal(x, y)
  raise [x, y].inspect unless x == y
end

class MaxStack
  def initialize
    @items = []
    @maxes = []
  end

  def push(value)
    if @maxes.empty? || value > @maxes[-1]
      @maxes.push(value)
    end
    @items.push(value)
  end

  def pop
    result = @items.pop
    @maxes.pop if result == @maxes[-1]
    result
  end

  # time: O(1)
  # space: O(1)
  def max
    @maxes[-1]
  end
end

stack = MaxStack.new
stack.push(1)
stack.push(2)
stack.push(3)
stack.push(2)

assert_equal 3, stack.max

stack.pop
stack.pop
assert_equal 2, stack.max

puts 'Yay!'