advent12/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/14/advent-of-code-2022-day-12/
   2
   3 import AoC
   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 ((<|), (|>), (><))
   9 import Data.Sequence ((|>))
  10 import Data.Foldable (foldl')
  11 import Data.Char
  12 import Control.Monad.Reader
  13 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  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 Position = V2 Int -- r, c
  22 type Grid = Array Position Int
  23
  24 data Mountain = Mountain
  25   { _grid :: Grid
  26   , _start :: Position
  27   , _goal :: Position
  28   } deriving (Eq, Ord, Show)
  29 makeLenses ''Mountain
  30
  31 type MountainContext = Reader Mountain
  32
  33 data Agendum  =
  34     Agendum { _current :: Position
  35             , _trail :: Q.Seq Position
  36             , _trailCost :: Int
  37             , _cost :: Int
  38             } deriving (Show, Eq)
  39 makeLenses ''Agendum
  40
  41 type Agenda = P.MinPQueue Int Agendum
  42
  43 type ExploredStates = S.Set Position
  44
  45 main :: IO ()
  46 main =
  47   do  dataFileName <- getDataFileName
  48       text <- readFile dataFileName
  49       let mountain = mkMountain text
  50       -- print mountain
  51       print $ part1 mountain
  52       print $ part2 mountain
  53
  54 part1, part2 :: Mountain -> Int
  55 part1 mountain = maybe 0 _cost result
  56     where s = mountain ^. start
  57           result = runReader (searchMountain s) mountain
  58
  59 part2 mountain = minimum results
  60   where starts = possibleStarts mountain
  61         results = fmap (runSearch mountain) starts
  62
  63 runSearch :: Mountain -> Position -> Int
  64 runSearch mountain s = maybe maxCost _cost result
  65   where result = runReader (searchMountain s) mountain
  66         maxCost = length $ indices $ mountain ^. grid
  67
  68 possibleStarts :: Mountain -> [Position]
  69 possibleStarts mountain = map fst $ filter ((== 0) . snd)
  70                                   $ assocs $ mountain ^. grid
  71
  72 mkMountain :: String -> Mountain
  73 mkMountain text = Mountain { _grid = grid, _start = s, _goal = g }
  74   where rows = lines text
  75         r = length rows - 1
  76         c = (length $ head rows) - 1
  77         grid0 = listArray ((V2 0 0), (V2 r c)) $ map mkCell $ concat rows
  78         mkCell e = ord e - ord 'a'
  79         s = head [i | i <- range ((V2 0 0), (V2 r c)), grid0 ! i == (mkCell 'S')]
  80         g = head [i | i <- range ((V2 0 0), (V2 r c)), grid0 ! i == (mkCell 'E')]
  81         grid = grid0 // [(s, mkCell 'a'), (g, mkCell 'z')]
  82
  83 searchMountain :: Position -> MountainContext (Maybe Agendum)
  84 searchMountain startPos =
  85     do agenda <- initAgenda startPos
  86        aStar agenda S.empty
  87
  88 initAgenda :: Position -> MountainContext Agenda
  89 initAgenda pos =
  90     do c <- estimateCost pos
  91        return $ P.singleton c Agendum { _current = pos, _trail = Q.empty, _trailCost = 0, _cost = c}
  92
  93 aStar ::  Agenda -> ExploredStates -> MountainContext (Maybe Agendum)
  94 aStar agenda closed
  95     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
  96     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
  97     | P.null agenda = return Nothing
  98     | otherwise =
  99         do  let (_, currentAgendum) = P.findMin agenda
 100             let reached = currentAgendum ^. current
 101             nexts <- candidates currentAgendum closed
 102             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
 103             reachedGoal <- isGoal reached
 104             if reachedGoal
 105             then return (Just currentAgendum)
 106             else if reached `S.member` closed
 107                  then aStar (P.deleteMin agenda) closed
 108                  else aStar newAgenda (S.insert reached closed)
 109
 110 candidates ::  Agendum -> ExploredStates -> MountainContext (Q.Seq Agendum)
 111 candidates agendum closed =
 112     do  let candidate = agendum ^. current
 113         let previous = agendum ^. trail
 114         let prevCost = agendum ^. trailCost
 115         succs <- successors candidate
 116         let nonloops = Q.filter (\s -> s `S.notMember` closed) succs
 117         mapM (makeAgendum previous prevCost) nonloops
 118
 119
 120 makeAgendum ::  Q.Seq Position -> Int -> Position -> MountainContext Agendum
 121 makeAgendum previous prevCost newPosition =
 122     do predicted <- estimateCost newPosition
 123        grid <- asks _grid
 124        let newTrail = previous |> newPosition
 125        let incurred = prevCost + 1
 126        return Agendum { _current = newPosition
 127                       , _trail = newTrail
 128                       , _trailCost = incurred
 129                       , _cost = incurred + predicted
 130                       }
 131
 132 isGoal :: Position -> MountainContext Bool
 133 isGoal here =
 134   do goal <- asks _goal
 135      return $ here == goal
 136
 137 successors :: Position -> MountainContext (Q.Seq Position)
 138 successors here =
 139   do grid <- asks _grid
 140      let heightHere = grid ! here
 141      let neighbours =
 142           filter (\p -> (grid ! p) - heightHere <= 1)
 143           $
 144           filter (inRange (bounds grid))
 145             [ here ^+^ delta
 146             | delta <- [V2 -1 0, V2 1 0, V2 0 -1, V2 0 1]
 147             ]
 148      let succs = Q.fromList neighbours
 149      return succs
 150
 151 estimateCost :: Position -> MountainContext Int
 152 estimateCost here =
 153   do goal <- asks _goal
 154      let (V2 dr dc) = here ^-^ goal
 155      return $ (abs dr) + (abs dc)