Day 9 - Part 2

src/Day9.hs

   ( parseMovements,
     Move (..),
     Grid (..),
-    startGrid,
     up,
     down,
     left,
     normalizeMovement,
     Step (..),
     follow,
-    day9
+    buildTails,
+    day9,
   )
 where
 
 import Control.Arrow ((>>>))
-import Control.Monad.Trans.State.Strict (State, runState, modify)
+import Control.Monad.Trans.State.Strict (State, modify, runState)
 import qualified Data.Set as S
 
 newtype Move = Move (Int, Int)
 
 data Grid = Grid {h :: Pos, t :: Pos} deriving (Show, Eq)
 
-data MovementLog = MovementLog { visited :: S.Set Pos, recordedSteps :: [Step] }
+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)
+logOne :: (Pos, Step) -> MovementLog -> MovementLog
+logOne (p, s) (MovementLog v steps) =
+  MovementLog (S.insert p v) (s : steps)
+
+buildTails :: [Step] -> [([Pos], [Step])]
+buildTails steps = iterate (snd >>> recordPositions) ([], steps)
 
 recordPositions :: [Step] -> ([Pos], [Step])
 recordPositions steps = (allPositions, allSteps)
   where
+    -- insert last tail position
     allPositions = S.toList $ S.insert lastTailPos (visited state)
     allSteps = reverse (recordedSteps state)
-    ((_, Grid { t = lastTailPos}), state) = runState (recordPositions' steps startGrid) (MovementLog S.empty [])
+    
+    ((_, Grid {t = lastTailPos}), state) =
+      runState (recordPositions' steps Grid {h = (0,0), t = (0,0)}) (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 modify (appendOne (newTail, followStep)) >> recordPositions' ms Grid {h = newHead, t = newTail}
-
-startGrid :: Grid
-startGrid = Grid {h = p, t = p}
-  where
-    p = (0, 0)
+       in modify (logOne (newTail, followStep))
+            >> recordPositions' ms Grid {h = newHead, t = newTail}
 
 parseMovements :: String -> [Move]
 parseMovements = lines >>> map toMove
 step :: Pos -> Step -> Pos
 step (x, y) (Step (dx, dy)) = (x + dx, y + dy)
 
-distance :: Pos -> Pos -> Int
-distance (x, y) (x2, y2) = floor (dx ** 2 + dy ** 2)
-  where
-    dx :: Double
-    dx = fromIntegral $ x - x2
-    dy :: Double
-    dy = fromIntegral $ y - y2
-
 follow :: Pos -> Pos -> Step
-follow p1@(x1, y1) p2@(x2, y2)
-  | distance p1 p2 > 2 = Step (signum dx, signum dy)
-  | abs dx == 2 = Step (signum dx, 0)
-  | abs dy == 2 = Step (0, signum dy)
+follow (x1, y1) (x2, y2)
+  | abs dx > 1 || abs dy > 1 = Step (signum dx, signum dy)
   | otherwise = Step (0, 0)
   where
     dx = x1 - x2
 
 day9 :: IO ()
 day9 = do
-   input <- readFile "ressources/day09-input"
-   putStrLn "Day9"
-   let movements = parseMovements input
-   let positions = concatMap normalizeMovement >>> recordPositions >>> fst >>> length $ movements
-   putStrLn ("Number of distinct positions " ++ show positions)
+  input <- readFile "ressources/day09-input"
+  putStrLn "Day9"
+  let headSteps = parseMovements >>> concatMap normalizeMovement $ input
+  let allTails = buildTails headSteps
+  let positions = (!! 1) >>> fst >>> length $ allTails
+  putStrLn ("Number of distinct positions " ++ show positions)
+  let tail9 = (!! 9) >>> fst >>> length $ allTails
+  putStrLn ("Number of distinct positions of tail 9" ++ show tail9)

test/Day9Spec.hs

   (0,0),(1,0),(2,0),(3,0)
   ]
 
+testInput2 :: String
+testInput2 = [str|R 5
+                 |U 8
+                 |L 8
+                 |D 3
+                 |R 17
+                 |D 10
+                 |L 25
+                 |U 20
+                 |]
+
 spec :: Spec
 spec =
   describe "Day9" $ do
         step (0, 1) (Step (0, -1)) `shouldBe` (0, 0)
       it "normalizes movement" $ do
         normalizeMovement still `shouldBe` []
-        normalizeMovement (left 1) `shouldBe` [Step (-1, 0)]
+        normalizeMovement (left 2) `shouldBe` [Step (-1, 0),Step (-1, 0)]
         normalizeMovement (right 1) `shouldBe` [Step (1, 0)]
         normalizeMovement (up 1) `shouldBe` [Step (0, 1)]
         normalizeMovement (down 2) `shouldBe` [Step (0, -1), Step (0, -1)]
         positions `shouldBe` expectedPositions
       it "solves the riddle" $ do
         input <- readFile "ressources/day09-input"
-        let movements = parseMovements input
-        let positions = concatMap normalizeMovement >>> recordPositions >>> fst >>> length $ movements
+        let headSteps = parseMovements >>> concatMap normalizeMovement $ input
+        let allTails = buildTails headSteps
+        let positions = (!! 1)  >>> fst >>> length $ allTails
         positions `shouldBe` 5878
+      it "solves example of part 2" $ do
+        let headSteps = parseMovements >>> concatMap normalizeMovement $ testInput2
+        let allTails = buildTails headSteps
+        let tail9Positions = fst $ allTails !! 9
+        length tail9Positions `shouldBe` 36
+      it "solves the riddle part 2" $ do
+        input <- readFile "ressources/day09-input"
+        let headSteps = parseMovements >>> concatMap normalizeMovement $ input
+        let allTails = buildTails headSteps
+        let positions = (!! 9)  >>> fst >>> length $ allTails
+        positions `shouldBe` 2405