1 -- Writeup at https://work.njae.me.uk/2023/12/21/advent-of-code-2023-day-17/
5 import qualified Data.PQueue.Prio.Min as P
6 import qualified Data.Set as S
7 import qualified Data.Sequence as Q
8 -- import Data.Sequence ((|>), (><), Seq( (:|>) ) )
9 import Data.Sequence ((|>), (><))
10 import Data.Foldable (foldl')
12 import Control.Monad.Reader
13 -- import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
14 import Control.Lens hiding ((|>))
15 import Linear (V2(..), (^+^), (^-^), (^*), _x, _y)
16 import Data.Array.IArray
18 type Position = V2 Int -- r, c
19 _r :: Lens' (V2 Int) Int
21 _c :: Lens' (V2 Int) Int
24 data Direction = U | D | L | R deriving (Show, Eq, Ord)
25 data Move = Move Direction Int deriving (Show, Eq, Ord)
27 type Trail = Q.Seq Move
29 type DirectedPosition = (Direction, Position)
31 type Grid = Array Position Int
33 type ExploredStates = S.Set DirectedPosition
39 } deriving (Eq, Ord, Show)
42 type CityContext = Reader City
48 Agendum { _current :: DirectedPosition
55 type Agenda a = P.MinPQueue Int (Agendum a)
59 do dataFileName <- getDataFileName
60 text <- readFile dataFileName
61 let city = mkCity text
66 part1, part2 :: City -> Int
67 part1 city = maybe 0 _cost result
68 where s = city ^. start
69 result = runReader (searchCity s) city :: (Maybe (Agendum Crucible))
71 part2 city = maybe 0 _cost result
72 where s = city ^. start
73 result = runReader (searchCity s) city :: (Maybe (Agendum UltraCrucible))
75 mkCity :: String -> City
76 mkCity text = City { _grid = grid, _start = (V2 0 0), _goal = (V2 r c) }
77 where rows = lines text
79 c = (length $ head rows) - 1
80 grid = listArray ((V2 0 0), (V2 r c)) $ map digitToInt $ concat rows
83 class Searchable a where
85 searchCity :: Position -> CityContext (Maybe (Agendum a))
87 do agenda <- initAgenda startPos
90 initAgenda :: Position -> CityContext (Agenda a)
92 do c <- estimateCost pos
93 let dAgendum = Agendum { _current = (D, pos), _trail = Q.empty, _trailCost = 0, _cost = c}
94 dNexts <- candidates dAgendum S.empty
95 let rAgendum = Agendum { _current = (R, pos), _trail = Q.empty, _trailCost = 0, _cost = c}
96 rNexts <- candidates rAgendum S.empty
97 let nexts = dNexts >< rNexts
98 let agenda = foldl' (\q a -> P.insert (_cost a) a q) P.empty nexts
101 aStar :: (Agenda a) -> ExploredStates -> CityContext (Maybe (Agendum a))
103 -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
104 -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
105 -- | trace ("Peeping " ++ (show $ snd $ P.findMin agenda) ) False = undefined
106 -- | trace ("Peeping " ++ (show agenda) ) False = undefined
107 | P.null agenda = return Nothing
109 do let (_, currentAgendum) = P.findMin agenda
110 let reached = currentAgendum ^. current
111 nexts <- candidates currentAgendum closed
112 let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
113 reachedGoal <- isGoal reached
115 then return (Just currentAgendum)
116 else if reached `S.member` closed
117 then aStar (P.deleteMin agenda) closed
118 else aStar newAgenda (S.insert reached closed)
120 candidates :: (Agendum a) -> ExploredStates -> CityContext (Q.Seq (Agendum a))
121 candidates agendum closed =
122 do let candidate = agendum ^. current
123 let (_, here) = candidate
124 let previous = agendum ^. trail
125 let prevCost = agendum ^. trailCost
126 succs <- successors agendum candidate
127 let nonloops = Q.filter (\s -> (endingDirPos here s) `S.notMember` closed) succs
128 mapM (makeAgendum previous prevCost here) nonloops
130 successors :: (Agendum a) -> DirectedPosition -> CityContext (Q.Seq Move)
132 makeAgendum :: Trail -> Int -> Position -> Move -> CityContext (Agendum a)
133 makeAgendum previous prevCost here move =
134 do let positions = toPositions here move
135 predicted <- estimateCost $ last positions
137 let newTrail = previous |> move
138 let incurred = prevCost + (sum $ fmap (grid !) positions)
139 return Agendum { _current = endingDirPos here move
141 , _trailCost = incurred
142 , _cost = incurred + predicted
146 instance Searchable Crucible where
147 successors _ = successorsWithRange (1, 3)
149 instance Searchable UltraCrucible where
150 successors _ = successorsWithRange (4, 10)
152 successorsWithRange :: (Int, Int) -> DirectedPosition -> CityContext (Q.Seq Move)
153 successorsWithRange rng (dir, here) =
154 do grid <- asks _grid
155 let moves = [ Move d n
156 | d <- turnDirections dir
159 let validMoves = filter (allInBounds (bounds grid) here) moves
160 return $ Q.fromList validMoves
162 isGoal :: DirectedPosition -> CityContext Bool
164 do goal <- asks _goal
165 return $ here == goal
167 estimateCost :: Position -> CityContext Int
169 do goal <- asks _goal
170 let (V2 dr dc) = here ^-^ goal
171 return $ (abs dr) + (abs dc)
173 delta :: Direction -> Position
179 turnDirections :: Direction -> [Direction]
180 turnDirections U = [L, R]
181 turnDirections D = [L, R]
182 turnDirections L = [U, D]
183 turnDirections R = [U, D]
185 toPositions :: Position -> Move -> [Position]
186 toPositions here (Move dir n) = [ here ^+^ (d ^* i) | i <- [1..n] ]
189 endingDirPos :: Position -> Move -> DirectedPosition
190 endingDirPos here move@(Move dir _) = (dir, last $ toPositions here move)
192 allInBounds :: (Position, Position) -> Position -> Move -> Bool
193 allInBounds bounds here move = all (inRange bounds) $ toPositions here move