[hs] Solve 2020_15 part 2

hs/src/Aoc/Y2020/D15.hs

@@ -1,8 +1,11 @@
+{-# LANGUAGE NumericUnderscores #-}
+{-# LANGUAGE Strict             #-}
+
 module Aoc.Y2020.D15
   ( day
   ) where
 
-import           Data.List
+import qualified Data.Map  as Map
 
 import           Aoc.Day
 import           Aoc.Parse
@@ -10,19 +13,65 @@ import           Aoc.Parse
 parser :: Parser [Int]
 parser = (decimal `sepBy` char ',') <* newline
 
-step :: [Int] -> [Int]
+data State = State
-step []     = error "list must not be empty"
+  { sIdx      :: Int
-step (x:xs) = case elemIndex x xs of
+  , sNum      :: Int
-  Nothing -> 0 : x : xs
+  , sLastSeen :: Map.Map Int Int
-  Just i  -> (i + 1) : x : xs
+  } deriving (Show)
 
-stepUntil :: Int -> [Int] -> [Int]
+-- nums must not be empty
-stepUntil amount nums = iterate step nums !! (amount - length nums)
+newState :: [Int] -> State
+newState nums = State
+  { sIdx = length nums - 1
+  , sNum = last nums
+  , sLastSeen = Map.fromList $ zip (init nums) [0..]
+  }
+
+step :: State -> State
+step s =
+  let newNum = case sLastSeen s Map.!? sNum s of
+        Nothing  -> 0
+        Just idx -> sIdx s - idx
+  in  State
+      { sIdx = sIdx s + 1
+      , sNum = newNum
+      , sLastSeen = Map.insert (sNum s) (sIdx s) $ sLastSeen s
+      }
+
+replicateF :: Int -> (a -> a) -> a -> a
+replicateF n f a
+  | n <= 0    = a
+  | otherwise = replicateF (n - 1) f (f a)
+
+stepUntil :: Int -> State -> IO State
+stepUntil amount s = stepUntilIo (amount - sIdx s - 1) s
+
+-- | A Helper for stepUntil that prints how many steps are left in regular
+-- intervals. Don't call it directly.
+stepUntilIo :: Int -> State -> IO State
+stepUntilIo amount s | amount <= 0 = pure s
+stepUntilIo amount s = do
+  putStrLn $ show amount ++ " left"
+  let width = min amount 100_000
+  stepUntilIo (amount - width) $ replicateF width step s
 
 solver :: [Int] -> IO ()
 solver nums = do
+  let s = newState nums
+
   putStrLn ">> Part 1"
-  print $ head $ stepUntil 2020 $ reverse nums
+  p1 <- sNum <$> stepUntil 2020 s
+  print p1
+
+  putStrLn ""
+  putStrLn ">> Part 2"
+  p2 <- sNum <$> stepUntil 30_000_000 s
+  print p2
+
+  putStrLn ""
+  putStrLn ">> Summary"
+  putStrLn $ "Part 1: " ++ show p1
+  putStrLn $ "Part 2: " ++ show p2
 
 day :: Day
 day = dayParse parser solver