advent17/MainSteps.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 Debug.Trace
   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 ((|>), Seq( (:|>) ) )
  11 import Data.Foldable (foldl', toList)
  12 import Data.Char
  13 import Control.Monad.Reader
  14 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  15 import Linear (V2(..), (^+^), (^-^), _x, _y)
  16 import Data.Array.IArray
  17
  18 type Position = V2 Int -- r, c
  19 _r :: Lens' (V2 Int) Int
  20 _r = _x
  21 _c :: Lens' (V2 Int) Int
  22 _c = _y
  23
  24 type Trail = Q.Seq Position
  25
  26 type Grid = Array Position Int
  27
  28 data City = City
  29   { _grid :: Grid
  30   , _start :: Position
  31   , _goal :: Position
  32   } deriving (Eq, Ord, Show)
  33 makeLenses ''City
  34
  35 type CityContext = Reader City
  36
  37 data Agendum  =
  38     Agendum { _current :: Trail
  39             , _trail :: Trail
  40             , _trailCost :: Int
  41             , _cost :: Int
  42             } deriving (Show, Eq)
  43 makeLenses ''Agendum
  44
  45 type Agenda = P.MinPQueue Int Agendum
  46
  47 type ExploredStates = S.Set Trail
  48
  49 main :: IO ()
  50 main =
  51   do  dataFileName <- getDataFileName
  52       text <- readFile dataFileName
  53       let city = mkCity text
  54       -- print city
  55       print $ part1 city
  56       -- print $ part2 city
  57
  58 -- part1, part2 :: City -> Int
  59 part1 city = maybe 0 _cost result
  60     where s = city ^. start
  61           result = runReader (searchCity s) city
  62
  63 -- part2 city = minimum results
  64 --   where starts = possibleStarts city
  65 --         results = fmap (runSearch city) starts
  66
  67 runSearch :: City -> Position -> Int
  68 runSearch city s = maybe maxCost _cost result
  69   where result = runReader (searchCity s) city
  70         maxCost = length $ indices $ city ^. grid
  71
  72
  73 mkCity :: String -> City
  74 mkCity text = City { _grid = grid, _start = (V2 0 0), _goal = (V2 r c) }
  75   where rows = lines text
  76         r = length rows - 1
  77         c = (length $ head rows) - 1
  78         grid = listArray ((V2 0 0), (V2 r c)) $ map digitToInt $ concat rows
  79
  80 searchCity :: Position -> CityContext (Maybe Agendum)
  81 searchCity startPos =
  82     do agenda <- initAgenda startPos
  83        aStar agenda S.empty
  84
  85 initAgenda :: Position -> CityContext Agenda
  86 initAgenda pos =
  87     do c <- estimateCost pos
  88       --  return $ P.singleton c Agendum { _current = pos, _trail = Q.empty, _trailCost = 0, _cost = c}
  89        return $ P.singleton c Agendum { _current = Q.singleton pos, _trail = Q.singleton pos, _trailCost = 0, _cost = c}
  90
  91 aStar ::  Agenda -> ExploredStates -> CityContext (Maybe Agendum)
  92 aStar agenda closed
  93     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
  94     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
  95     -- | trace ("Peeping " ++ (show $ snd $ P.findMin agenda) ) False = undefined
  96     -- | trace ("Peeping " ++ (show 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 -> CityContext (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 bent = Q.filter isBent succs
 117         let nonloops = Q.filter (\s -> s `S.notMember` closed) bent
 118         mapM (makeAgendum previous prevCost) nonloops
 119
 120 isBent :: Trail -> Bool
 121 -- isBent previous (V2 row col)
 122 --   | Q.length previous <= 3 = True
 123 --   | otherwise = not $
 124 --       all (\p -> p ^. _r == row) previous || all (\p -> p ^. _c == col) previous
 125 isBent trail
 126   | Q.length trail <= 4 = True
 127   | otherwise = not $ all id $ toList $ Q.zipWith (==) diffs $ Q.drop 1 diffs
 128   where diffs = Q.zipWith (^-^) trail $ Q.drop 1 trail
 129
 130
 131 makeAgendum ::  Trail -> Int -> Trail -> CityContext Agendum
 132 makeAgendum previous prevCost newState =
 133     do let (_ :|> newPosition) = newState
 134        predicted <- estimateCost newPosition
 135        grid <- asks _grid
 136        let newTrail = previous |> newPosition
 137        let incurred = prevCost + (grid ! newPosition)
 138        return Agendum { _current = newState
 139                       , _trail = newTrail
 140                       , _trailCost = incurred
 141                       , _cost = incurred + predicted
 142                       }
 143
 144 isGoal :: Trail -> CityContext Bool
 145 isGoal (_ :|> here) =
 146   do goal <- asks _goal
 147      return $ here == goal
 148
 149 successors :: Trail -> CityContext (Q.Seq Trail)
 150 successors trail@(ph :|> here) =
 151   do grid <- asks _grid
 152      let neighbours =
 153           filter (inRange (bounds grid))
 154             [ here ^+^ delta
 155             | delta <- [V2 -1 0, V2 1 0, V2 0 -1, V2 0 1]
 156             ]
 157      let neighbours' = if Q.null ph
 158                         then neighbours
 159                         else let (_ :|> ph') = ph
 160                              in filter (/= ph') neighbours
 161      let prev = takeL 4 trail
 162      let succs = Q.fromList $ fmap (prev :|>) neighbours'
 163      return succs
 164
 165 estimateCost :: Position -> CityContext Int
 166 -- estimateCost _ = return 0
 167 estimateCost here =
 168   do goal <- asks _goal
 169      let (V2 dr dc) = here ^-^ goal
 170      return $ (abs dr) + (abs dc)
 171
 172
 173 takeL :: Int -> Q.Seq a -> Q.Seq a
 174 takeL _ Q.Empty = Q.empty
 175 takeL 0 _ = Q.empty
 176 takeL n (xs :|> x) = (takeL (n-1) xs) :|> x