advent17/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/21/advent-of-code-2023-day-17/
   2
   3 import AoC
   4
   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')
  11 import Data.Char
  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
  17
  18 type Position = V2 Int -- r, c
  19 _r :: Lens' (V2 Int) Int
  20 _r = _x
  21 _c :: Lens' (V2 Int) Int
  22 _c = _y
  23
  24 data Direction = U | D | L | R deriving (Show, Eq, Ord)
  25 data Move = Move Direction Int deriving (Show, Eq, Ord)
  26
  27 type Trail = Q.Seq Move
  28
  29 type DirectedPosition = (Direction, Position)
  30
  31 type Grid = Array Position Int
  32
  33 type ExploredStates = S.Set DirectedPosition
  34
  35 data City = City
  36   { _grid :: Grid
  37   , _start :: Position
  38   , _goal :: Position
  39   } deriving (Eq, Ord, Show)
  40 makeLenses ''City
  41
  42 type CityContext = Reader City
  43
  44 data Crucible
  45 data UltraCrucible
  46
  47 data Agendum a =
  48     Agendum { _current :: DirectedPosition
  49             , _trail :: Trail
  50             , _trailCost :: Int
  51             , _cost :: Int
  52             } deriving (Show, Eq)
  53 makeLenses ''Agendum
  54
  55 type Agenda a = P.MinPQueue Int (Agendum a)
  56
  57 main :: IO ()
  58 main =
  59   do  dataFileName <- getDataFileName
  60       text <- readFile dataFileName
  61       let city = mkCity text
  62       -- print city
  63       print $ part1 city
  64       print $ part2 city
  65
  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))
  70
  71 part2 city = maybe 0 _cost result
  72     where s = city ^. start
  73           result = runReader (searchCity s) city :: (Maybe (Agendum UltraCrucible))
  74
  75 mkCity :: String -> City
  76 mkCity text = City { _grid = grid, _start = (V2 0 0), _goal = (V2 r c) }
  77   where rows = lines text
  78         r = length rows - 1
  79         c = (length $ head rows) - 1
  80         grid = listArray ((V2 0 0), (V2 r c)) $ map digitToInt $ concat rows
  81
  82
  83 class Searchable a where
  84
  85   searchCity :: Position -> CityContext (Maybe (Agendum a))
  86   searchCity startPos =
  87     do agenda <- initAgenda startPos
  88        aStar agenda S.empty
  89
  90   initAgenda :: Position -> CityContext (Agenda a)
  91   initAgenda pos =
  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
  99         return agenda
 100
 101   aStar ::  (Agenda a) -> ExploredStates -> CityContext (Maybe (Agendum a))
 102   aStar agenda closed
 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
 108       | otherwise =
 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
 114               if reachedGoal
 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)
 119
 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
 129
 130   successors :: (Agendum a) -> DirectedPosition -> CityContext (Q.Seq Move)
 131
 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
 136         grid <- asks _grid
 137         let newTrail = previous |> move
 138         let incurred = prevCost + (sum $ fmap (grid !) positions)
 139         return Agendum { _current = endingDirPos here move
 140                         , _trail = newTrail
 141                         , _trailCost = incurred
 142                         , _cost = incurred + predicted
 143                         }
 144
 145
 146 instance Searchable Crucible where
 147   successors _ = successorsWithRange (1, 3)
 148
 149 instance Searchable UltraCrucible where
 150   successors _ = successorsWithRange (4, 10)
 151
 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
 157                   , n <- range rng
 158                   ]
 159       let validMoves = filter (allInBounds (bounds grid) here) moves
 160       return $ Q.fromList validMoves
 161
 162 isGoal :: DirectedPosition -> CityContext Bool
 163 isGoal (_, here) =
 164   do goal <- asks _goal
 165      return $ here == goal
 166
 167 estimateCost :: Position -> CityContext Int
 168 estimateCost here =
 169   do goal <- asks _goal
 170      let (V2 dr dc) = here ^-^ goal
 171      return $ (abs dr) + (abs dc)
 172
 173 delta :: Direction -> Position
 174 delta U = V2 (-1) 0
 175 delta D = V2 1 0
 176 delta L = V2 0 (-1)
 177 delta R = V2 0 1
 178
 179 turnDirections :: Direction -> [Direction]
 180 turnDirections U = [L, R]
 181 turnDirections D = [L, R]
 182 turnDirections L = [U, D]
 183 turnDirections R = [U, D]
 184
 185 toPositions :: Position -> Move -> [Position]
 186 toPositions here (Move dir n) = [ here ^+^ (d ^* i) | i <- [1..n] ]
 187   where d = delta dir
 188
 189 endingDirPos :: Position -> Move -> DirectedPosition
 190 endingDirPos here move@(Move dir _) = (dir, last $ toPositions here move)
 191
 192 allInBounds :: (Position, Position) -> Position -> Move -> Bool
 193 allInBounds bounds here move = all (inRange bounds) $ toPositions here move