advent23/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2024/01/02/advent-of-code-2023-day-23/
   2
   3 import qualified Debug.Trace as DT
   4
   5 import AoC
   6 import Linear -- (V2(..), (^+^))
   7 import qualified Data.Set as S
   8 import qualified Data.Map.Strict as M
   9 import qualified Data.Sequence as Q
  10 import Data.Sequence (Seq( (:|>), (:<|) ) )
  11 import Control.Lens
  12 import Data.List (foldl')
  13 import Control.Monad.Reader
  14 import qualified Data.PQueue.Prio.Max as P
  15 import Data.Foldable
  16 import Data.Maybe
  17
  18 data Slide = SlideLeft | SlideRight | SlideUp | SlideDown
  19   deriving (Show, Eq)
  20
  21 type Position = V2 Int -- r, c
  22
  23 _r, _c :: Lens' (V2 Int) Int
  24 _r = _x
  25 _c = _y
  26
  27 type Grid = S.Set Position
  28 type Slides = M.Map Position Slide
  29
  30 data CompressedPath = CPath { _nextPos :: Position, _pathLen :: Int }
  31   deriving (Show, Eq)
  32 makeLenses ''CompressedPath
  33
  34 type CompressedMap = M.Map Position [CompressedPath]
  35 data Mountain = Mountain
  36   { _paths :: CompressedMap
  37   , _start :: Position
  38   , _goal :: Position
  39   } deriving (Eq, Show)
  40 makeLenses ''Mountain
  41
  42 type MountainContext = Reader Mountain
  43
  44 data Agendum =
  45     Agendum { _current :: Position
  46             , _trail :: Q.Seq Position
  47             , _trailCost :: Int
  48             , _cost :: Int
  49             } deriving (Show, Eq)
  50 makeLenses ''Agendum
  51
  52 type Agenda = P.MaxPQueue Int Agendum
  53
  54 type ExploredStates = M.Map Position Int
  55
  56
  57 main :: IO ()
  58 main =
  59   do  dataFileName <- getDataFileName
  60       text <- readFile dataFileName
  61       let (forest, slides, start, end) = mkGrid text
  62       -- print $ compress slides forest start end
  63       print $ part1 slides forest start end
  64       print $ part2 slides forest start end
  65
  66 part1, part2 :: Slides -> Grid -> Position -> Position -> Int
  67 part1 slides forest start end = searchCompressed $ Mountain cMap start end
  68   where cMap = compress slides forest start end
  69 part2 slides forest start end = searchCompressed $ Mountain cMap start end
  70   where cMap = compress M.empty forest start end
  71 -- part2 _ forest start end = maximum $ fmap (pathLength cMap) paths
  72 --   where cMap = compress M.empty forest start end
  73 --         paths = searchCompressed cMap start end
  74
  75 adjacents :: Position -> Slides -> Grid -> [Position]
  76 adjacents here slides forest = filter (`S.notMember` forest) $ fmap (here ^+^) deltas
  77   where deltas = case M.lookup here slides of
  78                   Nothing ->  [ V2 0 1, V2 1 0, V2 0 (-1), V2 (-1) 0 ]
  79                   Just SlideLeft -> [ V2 0 (-1) ]
  80                   Just SlideRight -> [ V2 0 1 ]
  81                   Just SlideUp -> [ V2 (-1) 0 ]
  82                   Just SlideDown -> [ V2 1 0 ]
  83
  84 searchStep :: Slides -> Grid -> [Position] -> [[Position]]
  85 searchStep _ _ [] = []
  86 searchStep slides forest path@(here:rest) = fmap (:path) valids
  87   where nexts = adjacents here slides forest
  88         valids = filter (`notElem` rest) nexts
  89
  90 search :: Slides -> Grid -> [Position] -> CompressedMap -> [[Position]] -> CompressedMap
  91 search _ _ _ foundPaths [] = foundPaths
  92 search slides forest goals foundPaths (current:agenda)
  93   | head current `elem` goals = search slides forest goals foundPaths' agenda
  94   | otherwise = search slides forest goals foundPaths (agenda ++ extendeds)
  95   where extendeds = searchStep slides forest current
  96         origin = last current
  97         foundPaths' = if origin == head current then foundPaths
  98                         else M.adjust (cp :) origin foundPaths
  99         cp = CPath (head current) (length current - 1)
 100
 101 -- collapsing the map
 102
 103 interestingPoints :: Slides -> Grid -> Position -> Position -> CompressedMap
 104 interestingPoints slides forest start end = M.fromList [(p, []) | p <- pointsSE]
 105   where Just minR = minimumOf (folded . _r) forest
 106         Just maxR = maximumOf (folded . _r) forest
 107         Just minC = minimumOf (folded . _c) forest
 108         Just maxC = maximumOf (folded . _c) forest
 109         points = [ V2 r c | r <- [(minR + 2)..(maxR - 2)]
 110                  , c <- [(minC + 1)..(maxC - 1)]
 111                  , (V2 r c) `S.notMember` forest
 112                  , (length $ adjacents (V2 r c) slides forest) > 2
 113                  ]
 114         pointsSE = start : end : points
 115
 116 compress :: Slides -> Grid -> Position -> Position -> CompressedMap
 117 compress slides forest start end = foldl' go compressed0 iPoints
 118   where compressed0 = interestingPoints slides forest start end
 119         iPoints = M.keys compressed0
 120         go com here = search slides forest iPoints com $ fmap (: [here]) $ adjacents here slides forest
 121
 122
 123 -- searchCompressed :: CompressedMap -> Position -> [[Position]] -> [[Position]] -> [[Position]]
 124 -- -- searchCompressed _ _ _ (c:_) _ | DT.trace (show c) False = undefined
 125 -- searchCompressed _ _ found [] = found
 126 -- searchCompressed map goal found (current:agenda)
 127 --   | head current == goal = searchCompressed map goal (current:found) agenda
 128 --   | otherwise = searchCompressed map goal found (nextPositions ++ agenda)
 129 --   where neighbours0 = map M.! (head current)
 130 --         neighbours = neighbours0 ^.. folded . filtered ((`notElem` current) . _nextPos)
 131 --         nextPositions = fmap ((: current) . _nextPos) neighbours
 132
 133 searchCompressed :: Mountain -> Int
 134 searchCompressed mountain = maybe 0 _trailCost result
 135   where result = runReader searchMountain mountain
 136
 137 searchMountain :: MountainContext (Maybe Agendum)
 138 searchMountain =
 139   do agenda <- initAgenda
 140      aStar agenda Nothing
 141
 142 initAgenda :: MountainContext Agenda
 143 initAgenda =
 144    do s <- asks _start
 145       c <- estimateCost Q.Empty s
 146       let agendum = Agendum { _current = s, _trail = Q.empty, _trailCost = 0, _cost = c}
 147       let agenda = P.singleton c agendum
 148       return agenda
 149
 150 aStar ::  Agenda -> (Maybe Agendum) -> MountainContext (Maybe Agendum)
 151 aStar agenda best
 152     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 153     -- | DT.trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ) False = undefined
 154     -- | DT.trace ("Peeping " ++ (show $ snd $ P.findMax agenda) ) False = undefined
 155     -- | DT.trace ("Peeping " ++ (show agenda) ) False = undefined
 156     | P.null agenda = return best
 157     | (fst $ P.findMax agenda) < maybe 0 _trailCost best = return best
 158     | otherwise =
 159         do  let (_, currentAgendum) = P.findMax agenda
 160             let reached = currentAgendum ^. current
 161             nexts <- candidates currentAgendum
 162             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMax agenda) nexts
 163             reachedGoal <- isGoal reached
 164             let best' = updateBest reachedGoal best currentAgendum
 165             -- let closed' = M.insert reached (currentAgendum ^. trailCost) closed
 166             if reachedGoal -- || (reached `S.member` closed)
 167                 then aStar (P.deleteMax agenda) best' -- closed'
 168                 else aStar newAgenda best' -- closed'
 169
 170 updateBest :: Bool -> Maybe Agendum -> Agendum -> Maybe Agendum
 171 updateBest False current _ = current
 172 updateBest True Nothing best
 173  -- | DT.trace ("Nothing " ++ show best) False = undefined
 174  | otherwise = Just best
 175 updateBest True (Just current) best
 176   -- | DT.trace (show current ++ " " ++ show best) False = undefined
 177   | (current ^. trailCost) > (best ^. trailCost) = Just current
 178   | otherwise = Just best
 179
 180
 181
 182
 183 candidates :: Agendum -> MountainContext (Q.Seq Agendum)
 184 candidates agendum =
 185   do  let here = agendum ^. current
 186       let previous = agendum ^. trail
 187       let prevCost = agendum ^. trailCost
 188       ts <- asks _paths
 189       let succs = Q.fromList $ ts M.! here
 190       -- succs <- successors candidate
 191       let nonloops = Q.filter (\s -> (s ^. nextPos) `notElem` previous) succs
 192       mapM (makeAgendum previous prevCost here) nonloops
 193
 194
 195 makeAgendum :: (Q.Seq Position) -> Int -> Position -> CompressedPath -> MountainContext Agendum
 196 makeAgendum previous prevCost here step =
 197    do let newTrail = previous :|> here
 198       predicted <- estimateCost newTrail $ step ^. nextPos
 199       -- ts <- asks _trails
 200       let incurred = prevCost + step ^. pathLen
 201       return Agendum { _current = step ^. nextPos
 202                      , _trail = newTrail
 203                      , _trailCost = incurred
 204                      , _cost = incurred + predicted
 205                      }
 206
 207
 208 isGoal :: Position -> MountainContext Bool
 209 isGoal here =
 210   do goal <- asks _goal
 211      return $ here == goal
 212
 213 estimateCost :: Q.Seq Position -> Position -> MountainContext Int
 214 estimateCost r e =
 215   do ts <- asks _paths
 216      let endCost = fromMaybe 0 $ maximumOf (folded . filtered ((`notElem` r) . _nextPos) . pathLen) $ ts M.! e
 217      let res = S.fromList $ toList (r :|> e)
 218      let otherPaths = concat $ M.elems $ ts `M.withoutKeys` res
 219      let restCost = sumOf (folded . filtered ((`notElem` r) . _nextPos) . pathLen) otherPaths
 220      return $ (restCost `div` 2) + endCost
 221
 222 -- pathLength :: CompressedMap -> [Position] -> Int
 223 -- pathLength map ps = sum $ zipWith (stepLength map) ps $ tail ps
 224
 225 -- stepLength :: CompressedMap -> Position -> Position -> Int
 226 -- stepLength map here there =
 227 --   -- head $ (map M.! there) ^.. folded . filtered ((== here) . _nextPos) . pathLen
 228 --   head $ (map M.! there) ^.. folded . filteredBy (nextPos . only here) . pathLen
 229
 230 -- reading the map
 231
 232 mkGrid :: String -> (Grid, Slides, Position, Position)
 233 mkGrid text = ((S.union forest caps), slides, start, end)
 234   where rows = lines text
 235         maxR = length rows - 1
 236         maxC = (length $ head rows) - 1
 237         forest = S.fromList [ V2 r c | r <- [0..maxR], c <- [0..maxC]
 238                             , rows !! r !! c == '#'
 239                             ]
 240         slides = M.fromList [ (V2 r c, readSlide (rows !! r !! c))
 241                             | r <- [0..maxR], c <- [0..maxC]
 242                             , elem (rows !! r !! c) ("<>^v" :: String)
 243                             ]
 244         start = head $ [ V2 0 c | c <- [0..maxC]
 245                        , rows !! 0 !! c == '.'
 246                        ]
 247         end = head $ [ V2 maxR c | c <- [0..maxC]
 248                      , rows !! maxR !! c == '.'
 249                      ]
 250         caps = S.fromList [start ^+^ (V2 -1 0), end ^+^ (V2 1 0)]
 251
 252 readSlide :: Char -> Slide
 253 readSlide '<' = SlideLeft
 254 readSlide '>' = SlideRight
 255 readSlide '^' = SlideUp
 256 readSlide 'v' = SlideDown