Keep track of each new max

1 files changed, 10 insertions, 2 deletions

diff --git a/2020/08/22/main.rb b/2020/08/22/main.rb
index 6042464..dc27bc7 100644
--- a/2020/08/22/main.rb
+++ b/2020/08/22/main.rb
@@ -5,18 +5,26 @@ 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
-    @items.pop
+    result = @items.pop
+    @maxes.pop if result == @maxes[-1]
+    result
   end
 
+  # time: O(1)
+  # space: O(1)
   def max
-    @items.max
+    @maxes[-1]
   end
 end