{-# LANGUAGE LambdaCase, MultiWayIf #-} module Main where import Data.Functor import Data.List import Data.Matrix hiding (trace) data Dir where N,E,W,S :: Dir deriving (Eq, Show) data Pos where Marked :: Pos Unmarked :: Pos Guard :: Dir -> Pos Crate :: Pos Obstacle :: Pos deriving Eq instance Show Pos where show = \case Unmarked -> "." Marked -> "X" Guard N -> "^" Guard E -> ">" Guard W -> "<" Guard S -> "v" Crate -> "#" Obstacle -> "O" parse :: String -> Matrix Pos parse = fromLists . (((\case '.' -> Unmarked 'X' -> Marked '^' -> Guard N '>' -> Guard E '<' -> Guard W 'v' -> Guard S '#' -> Crate 'O' -> Obstacle ) <$>) <$>) . lines guardPos :: Matrix Pos -> (Int, Int, Dir) guardPos s = let s' = toLists s Just y = findIndex (any (\case Guard _ -> True _ -> False)) s' Just x = findIndex (\case Guard _ -> True _ -> False) (s' !! y) d = (\case Guard a -> a _ -> error "" ) (getElem (y+1) (x+1) s) in (y+1,x+1,d) getMarked :: Matrix Pos -> [(Int,Int)] getMarked s = zip [1..] (toLists s) >>= \(y,l) -> zip [1..] l >>= \case (x,Marked) -> [(y,x)] (x,Guard _) -> [(y,x)] _ -> [] move :: Matrix Pos -> (Int, Int, Dir) -> Matrix Pos move s (y,x,d) = let (y',x',d') = case d of N -> (y-1,x,E) E -> (y,x+1,S) W -> (y,x-1,N) S -> (y+1,x,W) t = getElem y' x' s in case t of Crate -> setElem (Guard d') (y,x) s Obstacle -> setElem (Guard d') (y,x) s _ -> setElem (Guard d) (y',x') $ setElem Marked (y,x) s traversal :: Matrix Pos -> Matrix Pos traversal s = let i@(y,x,d) = guardPos s in if | y == 1 && d == N -> s | y == nrows s && d == S -> s | x == 1 && d == W -> s | x == ncols s && d == E -> s | otherwise -> traversal $ move s i solve1 :: Matrix Pos -> Int solve1 = sum . (length . filter (\case Marked -> True Guard _ -> True _ -> False ) <$>) . toLists . traversal isLoop :: Matrix Pos -> [(Int,Int,Dir)] -> Bool isLoop s v = let i@(y,x,d) = guardPos s s' = move s i in if | y == 1 && d == N -> False | y == nrows s && d == S -> False | x == 1 && d == W -> False | x == ncols s && d == E -> False | elem i v -> True | otherwise -> isLoop s' (i:v) solve2 :: Matrix Pos -> Int solve2 s = length . filter (`isLoop` []) . (<&> flip (setElem Obstacle) s) . (\\ [(\(a,b,_) -> (a,b)) $ guardPos s]) . getMarked . traversal $ s main :: IO () main = readFile "inputs/6" <&> parse >>= \i -> print (solve1 i) >> print (solve2 i)