day6 part two

Author: JasterV

src/Data/Matrix.hs

@@ -11,6 +11,7 @@ module Data.Matrix
     lookupMultiple,
     filterWithKey,
     size,
+    insert,
   )
 where
 
@@ -50,6 +51,13 @@ lookup position (Matrix hmap) = Map.lookup position hmap
 lookupMultiple :: [Position] -> Matrix v -> [v]
 lookupMultiple positions matrix = mapMaybe (`lookup` matrix) positions
 
+insert :: Position -> v -> Matrix v -> Matrix v
+insert position value (Matrix hmap) =
+  Matrix $
+    if Map.member position hmap
+      then Map.insert position value hmap
+      else hmap
+
 {--
   Search for the given value on the matrix.
   Return the position of the first match if found and nothing if it doens't exist.

src/Day6.hs

@@ -11,10 +11,24 @@ data Error = GuardNotFoundError
   deriving (Eq, Show)
 
 partOne :: String -> Either Error Int
-partOne input = length <$> predictGuardsRoute (parseLabMap input)
+partOne input = do
+  let labMap = parseLabMap input
+  guard <- findGuard labMap
+  let route = predictGuardsRoute guard labMap
+  return (length route)
 
 partTwo :: String -> Either Error Int
-partTwo _input = Right 0
+partTwo input = do
+  let labMap = parseLabMap input
+  guard <- findGuard labMap
+  let initialPosition = position guard
+      route = predictGuardsRoute guard labMap
+      candidates = filter (/= initialPosition) route
+
+  return $
+    length $
+      filter (hasLoop guard) $
+        map (\pos -> M.insert pos '#' labMap) candidates
 
 -- Laboratory Map
 type LabMap = Matrix Char
@@ -37,26 +51,33 @@ findGuard matrix =
         [_, _, _, Just pos] -> Right (Guard pos G.Left)
         _ -> Left GuardNotFoundError
 
-predictGuardsRoute :: LabMap -> Either Error [Position]
-predictGuardsRoute labMap = do
-  guard <- findGuard labMap
-  return $ go guard HashSet.empty HashSet.empty
+hasLoop :: Guard -> LabMap -> Bool
+hasLoop initialGuard labMap = go initialGuard HashSet.empty
+  where
+    go :: Guard -> Visited -> Bool
+    go guard visited =
+      let guard' = moveGuard guard labMap
+          visited' = HashSet.insert guard visited
+       in HashSet.member guard visited || ((guard /= guard') && go guard' visited')
+
+predictGuardsRoute :: Guard -> LabMap -> [Position]
+predictGuardsRoute initialGuard labMap = go initialGuard HashSet.empty HashSet.empty
   where
     go :: Guard -> Visited -> HashSet Position -> [Position]
     go guard visited acc =
-      let guard' = moveGuard guard
+      let guard' = moveGuard guard labMap
           acc' = HashSet.insert (position guard) acc
           visited' = HashSet.insert guard visited
        in -- If we have hit a loop or if the guard can't move anymore, finish prediction
-          if HashSet.member guard visited || (guard == guard')
+          if HashSet.member guard visited || guard == guard'
             then HashSet.toList acc'
             else go guard' visited' acc'
 
-    moveGuard :: Guard -> Guard
-    moveGuard guard =
-      let guard' = G.moveForward guard
-          mObstacle = M.lookup (position guard') labMap
-       in case mObstacle of
-            Nothing -> guard
-            Just '#' -> G.turnRight guard
-            Just _ -> guard'
+moveGuard :: Guard -> LabMap -> Guard
+moveGuard guard labMap =
+  let guard' = G.moveForward guard
+      mObstacle = M.lookup (position guard') labMap
+   in case mObstacle of
+        Nothing -> guard
+        Just '#' -> G.turnRight guard
+        Just _ -> guard'