38 lines
1 KiB
Haskell
38 lines
1 KiB
Haskell
import qualified Data.Map.Strict as Map
|
|
import Data.Map.Strict (Map)
|
|
import Debug.Trace
|
|
import Control.Monad.State
|
|
|
|
-- 1st is turn number
|
|
-- 2nd is map of spoken
|
|
-- 3rd is next number to speak
|
|
type S = (Int, Map Int Int, Int)
|
|
|
|
produceInitial lst =
|
|
let len = length lst
|
|
f acc (idx, el) = Map.insert el idx acc
|
|
initMap = foldl f Map.empty $ zip [0..] lst
|
|
in (len, initMap, 0)
|
|
|
|
speak :: Int -> State S Int
|
|
speak n = do
|
|
(turn, spoken, nextSpeak) <- get
|
|
let lastSpoke = Map.lookup nextSpeak spoken
|
|
let spoken2 = Map.insert nextSpeak turn spoken
|
|
let nextSpeak2 = case lastSpoke of
|
|
Nothing -> 0
|
|
Just n -> turn - n
|
|
let result = (turn + 1, spoken2, nextSpeak2)
|
|
-- modify (\_ -> traceShow result result)
|
|
modify (\_ -> result)
|
|
return nextSpeak
|
|
|
|
speakStream lst = (++) lst $ flip evalState initial $ mapM speak [0..]
|
|
where initial = produceInitial lst
|
|
|
|
actualInput = [5,1,9,18,13,8,0]
|
|
exInput = [0,3,6]
|
|
n = 30000000
|
|
-- n = 2020
|
|
|
|
main = do print (speakStream actualInput !! (n-1))
|