Day 9 - Part 2

src/Day9.hs

@@ -2,7 +2,6 @@ module Day9
   ( parseMovements,
     Move (..),
     Grid (..),
-    startGrid,
     up,
     down,
     left,
@@ -14,12 +13,13 @@ module Day9
     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)
@@ -50,31 +50,34 @@ type Pos = (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
@@ -93,19 +96,9 @@ normalizeMovement m = error $ "cannot normalize movement: " ++ show m
 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
@@ -113,8 +106,11 @@ follow p1@(x1, y1) p2@(x2, y2)
 
 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)