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