diff --git a/src/Day9.hs b/src/Day9.hs index 89ed3e1..522eab6 100644 --- a/src/Day9.hs +++ b/src/Day9.hs @@ -19,7 +19,8 @@ module Day9 where import Control.Arrow ((>>>)) -import Data.List (nub) +import Control.Monad.Trans.State.Strict (State, runState, modify) +import qualified Data.Set as S newtype Move = Move (Int, Int) deriving (Show, Eq) @@ -49,17 +50,26 @@ type Pos = (Int, Int) data Grid = Grid {h :: Pos, t :: Pos} deriving (Show, Eq) -recordPositions :: [Step] -> [Pos] -recordPositions steps = nub (lastTailPos:positions) +data MovementLog = MovementLog { visited :: S.Set Pos, recordedSteps :: [Step] } + deriving (Show, Eq) + +appendOne :: (Pos, Step) -> MovementLog -> MovementLog +appendOne (p, s) (MovementLog v steps) = + MovementLog (S.insert p v) (s:steps) + +recordPositions :: [Step] -> ([Pos], [Step]) +recordPositions steps = (allPositions, allSteps) where - (positions, _, Grid { t = lastTailPos}) = recordPositions' [] steps startGrid - recordPositions' :: [Pos] -> [Step] -> Grid -> ([Pos], [Step], Grid) - recordPositions' pos [] grid = (pos, [], grid) - recordPositions' pos (m : ms) Grid {h = headPos, t = tailPos} = + allPositions = S.toList $ S.insert lastTailPos (visited state) + allSteps = reverse (recordedSteps state) + ((_, Grid { t = lastTailPos}), state) = runState (recordPositions' steps startGrid) (MovementLog S.empty []) + recordPositions' :: [Step] -> Grid -> State MovementLog ([Step], Grid) + recordPositions' [] grid = return ([], grid) + recordPositions' (m : ms) Grid {h = headPos, t = tailPos} = let newHead = headPos `step` m followStep = follow newHead tailPos newTail = tailPos `step` followStep - in recordPositions' (tailPos : pos) ms Grid {h = newHead, t = newTail} + in modify (appendOne (newTail, followStep)) >> recordPositions' ms Grid {h = newHead, t = newTail} startGrid :: Grid startGrid = Grid {h = p, t = p} @@ -106,5 +116,5 @@ day9 = do input <- readFile "ressources/day09-input" putStrLn "Day9" let movements = parseMovements input - let positions = concatMap normalizeMovement >>> recordPositions >>> length $ movements + let positions = concatMap normalizeMovement >>> recordPositions >>> fst >>> length $ movements putStrLn ("Number of distinct positions " ++ show positions) diff --git a/test/Day9Spec.hs b/test/Day9Spec.hs index 4f4679e..41e855c 100644 --- a/test/Day9Spec.hs +++ b/test/Day9Spec.hs @@ -67,13 +67,13 @@ spec = let positions = parseMovements >>> concatMap normalizeMovement >>> recordPositions + >>> fst >>> sort $ testInput length positions `shouldBe` 13 positions `shouldBe` expectedPositions it "solves the riddle" $ do input <- readFile "ressources/day09-input" - putStrLn "Day9" let movements = parseMovements input - let positions = concatMap normalizeMovement >>> recordPositions >>> length $ movements + let positions = concatMap normalizeMovement >>> recordPositions >>> fst >>> length $ movements positions `shouldBe` 5878