advent24/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/24/advent-of-code-2022-day-24/
   2
   3 -- import Debug.Trace
   4
   5 import AoC
   6 import qualified Data.PQueue.Prio.Min as P
   7 import qualified Data.Set as S
   8 import qualified Data.IntMap.Strict as M
   9 import qualified Data.Sequence as Q
  10 -- import Data.Sequence ((<|), (|>), (><))
  11 import Data.Sequence ((|>))
  12 import Control.Monad.Reader
  13 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  14 import Linear (V2(..), (^+^), (^-^))
  15 import Data.Array.IArray
  16 -- import Data.Ix
  17 import Data.List
  18 import Data.Maybe
  19
  20 -- pattern Empty   <- (Q.viewl -> Q.EmptyL)  where Empty = Q.empty
  21 -- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<)  = (Q.<|)
  22 -- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>)  = (Q.|>)
  23
  24 type Position = V2 Int -- x, y
  25
  26 data Blizzard = Blizzard { _positionB :: Position, _headingB :: Position}
  27   deriving (Eq, Ord, Show)
  28 makeLenses ''Blizzard
  29
  30 type SafeValley = Array Position Bool
  31 type TimedValley = M.IntMap SafeValley
  32
  33 data Valley = Valley
  34   { blizzardStates :: TimedValley
  35   , start :: Position
  36   , goal :: Position
  37   } deriving (Eq, Ord, Show)
  38
  39 type ValleyContext = Reader Valley
  40
  41 data Explorer = Explorer
  42   { _currentPosition :: Position
  43   , _currentTime :: Int
  44   }deriving (Eq, Ord, Show)
  45 makeLenses ''Explorer
  46
  47 data Agendum  =
  48     Agendum { _current :: Explorer
  49             , _trail :: Q.Seq Explorer
  50             , _trailCost :: Int
  51             , _cost :: Int
  52             } deriving (Show, Eq)
  53 makeLenses ''Agendum
  54
  55 type Agenda = P.MinPQueue Int Agendum
  56
  57 type ExploredStates = S.Set Explorer
  58
  59 main :: IO ()
  60 main =
  61   do  dataFileName <- getDataFileName
  62       text <- readFile dataFileName
  63       let (blizzards, bnds) = mkInitialMap text
  64       let valley = makeValley bnds blizzards 1000
  65       print $ part1 valley
  66       print $ part2 valley
  67
  68 part1, part2 :: Valley -> Int
  69 part1 valley = _currentTime $ _current $ fromJust result
  70   where result = runSearch valley 0
  71
  72 part2 valley = trip3End
  73   where reverseValley = valley {start = (goal valley), goal = (start valley)}
  74         trip1End = _currentTime $ _current $ fromJust $ runSearch valley 0
  75         trip2End = _currentTime $ _current $ fromJust $ runSearch reverseValley trip1End
  76         trip3End = _currentTime $ _current $ fromJust $ runSearch valley trip2End
  77
  78 makeValley :: (Position, Position) -> S.Set Blizzard -> Int -> Valley
  79 makeValley bds blizzards n = Valley
  80   { blizzardStates = bStates
  81   , start = V2 (minX + 1) maxY
  82   , goal = V2 (maxX - 1) minY
  83   }
  84   where bStates = simulateBlizzards bds blizzards n
  85         (V2 minX minY, V2 maxX maxY) = bounds $ bStates M.! 0
  86
  87 runSearch :: Valley -> Int -> Maybe Agendum
  88 runSearch valley t = result
  89   where result = runReader (searchValley t) valley
  90
  91 searchValley :: Int -> ValleyContext (Maybe Agendum)
  92 searchValley t =
  93     do agenda <- initAgenda t
  94        aStar agenda S.empty
  95
  96 initAgenda :: Int -> ValleyContext Agenda
  97 initAgenda t =
  98     do pos <- asks start
  99        let explorer = Explorer pos t
 100        c <- estimateCost explorer
 101        return $ P.singleton c Agendum { _current = explorer, _trail = Q.empty, _trailCost = 0, _cost = c}
 102
 103 aStar ::  Agenda -> ExploredStates -> ValleyContext (Maybe Agendum)
 104 aStar agenda closed
 105     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 106     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
 107     -- | trace ("Peeping " ++ (show $ snd $ P.findMin agenda) ) False = undefined
 108     | P.null agenda = return Nothing
 109     | otherwise =
 110         do  let (_, currentAgendum) = P.findMin agenda
 111             let reached = currentAgendum ^. current
 112             nexts <- candidates currentAgendum closed
 113             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
 114             reachedGoal <- isGoal reached
 115             if reachedGoal
 116             then return (Just currentAgendum)
 117             else if reached `S.member` closed
 118                  then aStar (P.deleteMin agenda) closed
 119                  else aStar newAgenda (S.insert reached closed)
 120
 121 candidates ::  Agendum -> ExploredStates -> ValleyContext (Q.Seq Agendum)
 122 candidates agendum closed =
 123     do  let candidate = agendum ^. current
 124         let previous = agendum ^. trail
 125         let prevCost = agendum ^. trailCost
 126         succs <- successors candidate
 127         let nonloops = Q.filter (\s -> s `S.notMember` closed) succs
 128         mapM (makeAgendum previous prevCost) nonloops
 129
 130 makeAgendum ::  Q.Seq Explorer -> Int -> Explorer -> ValleyContext Agendum
 131 makeAgendum previous prevCost newExplorer =
 132     do predicted <- estimateCost newExplorer
 133        let newTrail = previous |> newExplorer
 134        let incurred = prevCost + 1
 135        return Agendum { _current = newExplorer
 136                       , _trail = newTrail
 137                       , _trailCost = incurred
 138                       , _cost = incurred + predicted
 139                       }
 140
 141 isGoal :: Explorer -> ValleyContext Bool
 142 isGoal here =
 143   do goal <- asks goal
 144      return $ (here ^. currentPosition) == goal
 145
 146 successors :: Explorer -> ValleyContext (Q.Seq Explorer)
 147 successors here =
 148   do allBlizzards <- asks blizzardStates
 149      let nextTime = (here ^. currentTime) + 1
 150      let blizzards = allBlizzards M.! nextTime
 151      let bds = bounds blizzards
 152      let pos = here ^. currentPosition
 153      let neighbours =
 154           filter (\p -> (blizzards ! p)) $
 155           filter (inRange bds)
 156             [ pos ^+^ delta
 157             | delta <- [V2 0 0, V2 -1 0, V2 1 0, V2 0 -1, V2 0 1]
 158             ]
 159      let succs = Q.fromList
 160                   $ fmap (\nbr -> here & currentTime .~ nextTime
 161                                        & currentPosition .~ nbr )
 162                          neighbours
 163      return succs
 164
 165 estimateCost :: Explorer -> ValleyContext Int
 166 estimateCost here =
 167   do goal <- asks goal
 168      let (V2 dx dy) = (here ^. currentPosition) ^-^ goal
 169      return $ (abs dx) + (abs dy)
 170
 171
 172 mkInitialMap :: String -> (S.Set Blizzard, (Position, Position))
 173 mkInitialMap text =
 174   ( S.fromList [ Blizzard (V2 (x - 1) (y - 1)) (deltaOfArrow $ charAt x y)
 175                | x <- [0..maxX]
 176                , y <- [0..maxY]
 177                , isBlizzard x y
 178                ]
 179   , (V2 0 0, V2 (maxX - 1) (maxY - 1))
 180   )
 181   where rows = reverse $ lines text
 182         maxY = length rows - 1
 183         maxX = (length $ head rows) - 1
 184         charAt x y = ((rows !! y) !! x)
 185         isBlizzard x y = (charAt x y) `elem` ("^<>v" :: String)
 186
 187 deltaOfArrow :: Char -> Position
 188 deltaOfArrow '^' = V2 0 1
 189 deltaOfArrow '>' = V2 1 0
 190 deltaOfArrow 'v' = V2 0 -1
 191 deltaOfArrow '<' = V2 -1 0
 192 deltaOfArrow _   = V2 0 0
 193
 194 advanceBlizzard :: (Position, Position) -> S.Set Blizzard -> S.Set Blizzard
 195 advanceBlizzard bnds blizzards = S.map (advanceOneBlizzard bnds) blizzards
 196
 197 advanceOneBlizzard :: (Position, Position) -> Blizzard -> Blizzard
 198 advanceOneBlizzard (_, V2 maxX maxY) blizzard = blizzard' & positionB %~ wrap
 199   where wrap (V2 x0 y0) = V2 (x0 `mod` maxX) (y0 `mod` maxY)
 200         blizzard' = blizzard & positionB %~ (^+^ (blizzard ^. headingB))
 201
 202 toSafe :: (Position, Position) -> S.Set Blizzard -> SafeValley
 203 toSafe (_, V2 maxX maxY) blizzards = accumArray (\_ _ -> False) True bnds' unsafeElements
 204   where unsafeElements = fmap (\i -> (i, False)) $ blizzardLocations ++ walls
 205         blizzardLocations = fmap (^+^ (V2 1 1)) $ fmap (^. positionB) $ S.toList blizzards
 206         walls = left ++ right ++ top ++ bottom
 207         left   = range (V2 0          0         , V2 0          (maxY + 1))
 208         right  = range (V2 (maxX + 1) 0         , V2 (maxX + 1) (maxY + 1))
 209         top    = range (V2 2          (maxY + 1), V2 (maxX + 1) (maxY + 1))
 210         bottom = range (V2 0          0         , V2 (maxX - 1) 0         )
 211         bnds' = (V2 0 0, V2 (maxX + 1) (maxY + 1))
 212
 213 simulateBlizzards :: (Position, Position) -> S.Set Blizzard -> Int -> TimedValley
 214 simulateBlizzards bnds blizzards n =
 215   M.fromList $ take n
 216              $ zip [0..]
 217              $ fmap (toSafe bnds)
 218              $ iterate (advanceBlizzard bnds) blizzards
 219
 220 showSafe :: SafeValley -> String
 221 showSafe valley = unlines $ reverse rows
 222   where (V2 minX minY, V2 maxX maxY) = bounds valley
 223         rows = [mkRow y | y <- [minY..maxY]]
 224         mkRow y = [if valley ! (V2 x y) then '.' else '#' | x <- [minX..maxX]]