advent15/Main.hs

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