diff --git a/src/Data/Matrix.hs b/src/Data/Matrix.hs index 3e56d4a..93264ce 100644 --- a/src/Data/Matrix.hs +++ b/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. diff --git a/src/Day6.hs b/src/Day6.hs index 395b1be..f46f7ca 100644 --- a/src/Day6.hs +++ b/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'