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