advent15/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/16/advent-of-code-2021-day-15/
   2
   3 -- import Debug.Trace
   4
   5 -- import qualified Data.Text.IO as TIO
   6
   7 import qualified Data.PQueue.Prio.Min as P
   8 import qualified Data.Set as S
   9 import qualified Data.Sequence as Q
  10 import Data.Sequence ((<|), (|>), (><)) --, ViewR( (:>) ), ViewL( (:<) ))
  11 import Data.Foldable (foldl', sum) -- (toList, foldr', foldl', all)
  12 import Data.Char
  13 import Control.Monad.Reader
  14 import Control.Lens hiding ((<|), (|>), (:>), (:<))
  15 -- import Data.Maybe (fromMaybe)
  16 import Linear (V2(..), (^+^), (^-^), (*^), (^*))
  17 import Data.Array.IArray
  18
  19 pattern Empty   <- (Q.viewl -> Q.EmptyL)  where Empty = Q.empty
  20 pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<)  = (Q.<|)
  21 pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>)  = (Q.|>)
  22
  23 type BasePosition = V2 Int -- r, c
  24 newtype Position = Position BasePosition -- r, c
  25   deriving (Eq, Ord, Show)
  26 newtype TiledPosition = TiledPosition BasePosition -- r, c
  27   deriving (Eq, Ord, Show)
  28 type Grid = Array BasePosition Int
  29
  30 data Cave = Cave
  31   { _grid :: Grid
  32   , _goal :: BasePosition
  33   } deriving (Eq, Ord, Show)
  34 makeLenses ''Cave
  35
  36 type CaveContext = Reader Cave
  37
  38 data Agendum s =
  39     Agendum { _current :: s
  40             , _trail :: Q.Seq s
  41             , _trailCost :: Int
  42             , _cost :: Int
  43             } deriving (Show, Eq)
  44 makeLenses ''Agendum
  45
  46 type Agenda s = P.MinPQueue Int (Agendum s)
  47
  48 type ExploredStates s = S.Set s
  49
  50 class (Eq s, Ord s, Show s) => SearchState s where
  51     unwrapPos :: s -> BasePosition
  52     emptySearchState :: s
  53     successors :: s -> CaveContext (Q.Seq s)
  54     estimateCost :: s -> CaveContext Int
  55     isGoal :: s -> CaveContext Bool
  56     entryCost :: s -> CaveContext Int
  57
  58
  59 instance SearchState Position where
  60
  61   unwrapPos (Position p) = p
  62
  63   emptySearchState = Position (V2 0 0)
  64
  65   -- successors :: Position -> CaveContext (Q.Seq Position)
  66   successors (Position here) =
  67     do grid <- asks _grid
  68        let neighbours =
  69             filter (inRange (bounds grid))
  70               [ here ^+^ delta
  71               | delta <- [V2 -1 0, V2 1 0, V2 0 -1, V2 0 1]
  72               ]
  73        let succs = Q.fromList $ map Position neighbours
  74        return succs
  75
  76   -- estimateCost :: Position -> CaveContext Int
  77   estimateCost (Position here) =
  78     do goal <- asks _goal
  79        let (V2 dr dc) = here ^-^ goal
  80        return $ (abs dr) + (abs dc)
  81
  82   -- isGoal :: here -> CaveContext Bool
  83   isGoal (Position here) =
  84     do goal <- asks _goal
  85        return $ here == goal
  86
  87   entryCost (Position here) =
  88     do grid <- asks _grid
  89        return $ grid ! here
  90
  91 instance SearchState TiledPosition where
  92
  93   unwrapPos (TiledPosition p) = p
  94
  95   emptySearchState = TiledPosition (V2 0 0)
  96
  97   -- successors :: Position -> CaveContext (Q.Seq Position)
  98   successors (TiledPosition here) =
  99     do grid <- asks _grid
 100        let (lowBound, highBound) = bounds grid
 101        let extendedBounds = ( lowBound
 102                             , tileScale highBound
 103                             )
 104        let neighbours =
 105             filter (inRange extendedBounds)
 106               [ here ^+^ delta
 107               | delta <- [V2 -1 0, V2 1 0, V2 0 -1, V2 0 1]
 108               ]
 109        let succs = Q.fromList $ map TiledPosition neighbours
 110        return succs
 111
 112   -- estimateCost :: Position -> CaveContext Int
 113   estimateCost (TiledPosition here) =
 114     do goal <- asks _goal
 115        let (V2 dr dc) = here ^-^ (tileScale goal)
 116        return $ (abs dr) + (abs dc)
 117
 118   -- isGoal :: here -> CaveContext Bool
 119   isGoal (TiledPosition here) =
 120     do goal <- asks _goal
 121        return $ here == (tileScale goal)
 122
 123   entryCost (TiledPosition (V2 r c)) =
 124     do grid <- asks _grid
 125        let (_, V2 maxR maxC) = bounds grid
 126        let (tileR, gridR) = r `divMod` (maxR + 1)
 127        let (tileC, gridC) = c `divMod` (maxC + 1)
 128        let gridCost = grid ! (V2 gridR gridC)
 129        let !cost = (gridCost - 1 + tileR + tileC) `mod` 9 + 1
 130        return cost
 131
 132 tileScale :: BasePosition -> BasePosition
 133 tileScale (V2 r c) = V2 (ts r) (ts c)
 134   where ts n = (n + 1) * 5 - 1
 135
 136 ------------------------------
 137
 138 main :: IO ()
 139 main =
 140   do  text <- readFile "data/advent15.txt"
 141       let cave = mkCave text
 142       print $ part1 cave
 143       print $ part2 cave
 144       -- print $ part2 grid
 145
 146 mkCave :: String -> Cave
 147 mkCave text = Cave { _grid = grid, _goal = V2 r c }
 148   where rows = lines text
 149         r = length rows - 1
 150         c = (length $ head rows) - 1
 151         grid = listArray ((V2 0 0), (V2 r c)) $ map mkCell $ concat rows
 152         mkCell e = digitToInt e
 153
 154
 155 part1 :: Cave -> Int
 156 part1 cave = maybe 0 _cost result
 157     where result = runReader searchCave cave :: Maybe (Agendum Position)
 158
 159 part2 :: Cave -> Int
 160 part2 cave = maybe 0 _cost result
 161     where result = runReader searchCave cave :: Maybe (Agendum TiledPosition)
 162
 163
 164 searchCave ::  SearchState s => CaveContext (Maybe (Agendum s))
 165 searchCave =
 166     do agenda <- initAgenda
 167        aStar agenda S.empty
 168
 169 initAgenda ::  SearchState s => CaveContext (Agenda s)
 170 initAgenda =
 171     do let ss = emptySearchState
 172        c <- estimateCost ss
 173        return $ P.singleton c Agendum { _current = ss, _trail = Q.empty, _trailCost = 0, _cost = c}
 174
 175
 176 aStar ::  SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
 177 aStar agenda closed
 178     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 179     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
 180     | P.null agenda = return Nothing
 181     | otherwise =
 182         do  let (_, currentAgendum) = P.findMin agenda
 183             let reached = currentAgendum ^. current
 184             nexts <- candidates currentAgendum closed
 185             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
 186             reachedGoal <- isGoal reached
 187             if reachedGoal
 188             then return (Just currentAgendum)
 189             else if reached `S.member` closed
 190                  then aStar (P.deleteMin agenda) closed
 191                  else aStar newAgenda (S.insert reached closed)
 192
 193
 194 candidates ::  SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
 195 candidates agendum closed =
 196     do  let candidate = agendum ^. current
 197         let previous = agendum ^. trail
 198         let prevCost = agendum ^. trailCost
 199         succs <- successors candidate
 200         let nonloops = Q.filter (\s -> s `S.notMember` closed) succs
 201         mapM (makeAgendum previous prevCost) nonloops
 202
 203
 204 makeAgendum ::  SearchState s => (Q.Seq s) -> Int -> s -> CaveContext (Agendum s)
 205 makeAgendum previous prevCost newPosition =
 206     do predicted <- estimateCost newPosition
 207        grid <- asks _grid
 208        let newTrail = previous |> newPosition
 209        newPositionCost <- entryCost newPosition
 210        let incurred = prevCost + newPositionCost
 211        return Agendum { _current = newPosition
 212                       , _trail = newTrail
 213                       , _trailCost = incurred
 214                       , _cost = incurred + predicted
 215                       }
 216