127 lines
3.7 KiB
Haskell
127 lines
3.7 KiB
Haskell
{-# LANGUAGE LambdaCase, MultiWayIf #-}
|
|
module Main where
|
|
|
|
import Data.Functor
|
|
import Data.List
|
|
import Data.Matrix hiding (trace)
|
|
import Debug.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)
|
|
|
|
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)
|
|
|
|
|
|
traversal' :: Matrix Pos -> [(Int,Int)] -> [(Int,Int)]
|
|
traversal' s v = let i@(y,x,d) = guardPos s
|
|
inf = case d of
|
|
N -> (y-1,x)
|
|
E -> (y,x+1)
|
|
W -> (y,x-1)
|
|
S -> (y+1,x)
|
|
in
|
|
if | y == 1 && d == N -> v
|
|
| y == nrows s && d == S -> v
|
|
| x == 1 && d == W -> v
|
|
| x == ncols s && d == E -> v
|
|
| otherwise ->
|
|
let t = (setElem Obstacle inf s) in
|
|
--trace ("t: " ++ show t) $
|
|
if isLoop t []
|
|
then traversal' (move s i) ((y,x):v)
|
|
else traversal' (move s i) v
|
|
|
|
|
|
solve2 :: Matrix Pos -> Int
|
|
solve2 s = length . (`traversal'` []) $ s
|
|
|
|
main :: IO ()
|
|
main = readFile "inputs/6" <&> parse >>= \i ->
|
|
print (solve1 i) >>
|
|
print (solve2 i)
|