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