diff --git a/5.hs b/5.hs index f2dc25b..38a63e5 100644 --- a/5.hs +++ b/5.hs @@ -1,24 +1,27 @@ -import Data.Sequence +import Data.IntMap (IntMap, insert, size, fromList, findWithDefault) +type Length = Int type Index = Int type Steps = Int -type State = (Steps, Index, Seq Int) +type State = (Steps, Index, IntMap Int) type Update = Int -> Int next :: Update -> State -> State next f (steps, i, jumps) = - let value = index jumps i + let value = findWithDefault undefined i jumps nextI = i + value - nextJumps = update i (f value) jumps + nextJumps = insert i (f value) jumps in (steps + 1, nextI, nextJumps) -getExitSteps :: Update -> State -> Int -getExitSteps f (steps, i, jumps) = - if i >= Data.Sequence.length jumps then steps else getExitSteps f $ next f (steps, i, jumps) +getExitSteps :: Length -> Update -> State -> Int +getExitSteps len f (steps, i, jumps) = + if i >= len then steps else getExitSteps len f $ next f (steps, i, jumps) main :: IO () main = do input <- readFile "5.txt" - let jumps = fromList . map read $ lines input :: Seq Int - print $ getExitSteps (+1) (0, 0, jumps) - print $ getExitSteps (\v -> if v >= 3 then v - 1 else v + 1) (0, 0, jumps) \ No newline at end of file + let jumpsList = map read $ lines input + jumpsMap = fromList $ zip [0..] jumpsList + len = length jumpsList + print $ getExitSteps len (+1) (0, 0, jumpsMap) + print $ getExitSteps len (\v -> if v >= 3 then v - 1 else v + 1) (0, 0, jumpsMap) \ No newline at end of file