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 Control.Lens
  10 import Data.List (foldl')
  11
  12 import Control.Monad.Par
  13 -- import Control.Monad.Par.Scheds.Trace
  14 -- import Control.Monad.Par.Scheds.Sparks
  15 import Control.Monad.Par.Combinator
  16
  17 data Slide = SlideLeft | SlideRight | SlideUp | SlideDown
  18   deriving (Show, Eq)
  19
  20 type Position = V2 Int -- r, c
  21
  22 _r, _c :: Lens' (V2 Int) Int
  23 _r = _x
  24 _c = _y
  25
  26 type Grid = S.Set Position
  27 type Slides = M.Map Position Slide
  28
  29 data CompressedPath = CPath { _nextPos :: Position, _pathLen :: Int }
  30   deriving (Show, Eq)
  31 makeLenses ''CompressedPath
  32
  33 type CompressedMap = M.Map Position [CompressedPath]
  34
  35
  36 parallelDepthLimit = 7 :: Int
  37
  38 main :: IO ()
  39 main =
  40   do  dataFileName <- getDataFileName
  41       text <- readFile dataFileName
  42       let (forest, slides, start, end) = mkGrid text
  43       print $ part1 slides forest start end
  44       print $ part2 slides forest start end
  45
  46 part1, part2 :: Slides -> Grid -> Position -> Position -> Int
  47 part1 slides forest start end = maximum $ fmap (pathLength cMap) paths
  48   where cMap = compress slides forest start end
  49         paths = searchCompressed cMap end [start]
  50 part2 _ forest start end = maximum $ fmap (pathLength cMap) paths
  51 -- part2 _ forest start end = maximum $ fmap length paths
  52   where cMap = compress M.empty forest start end
  53         paths = searchCompressed cMap end [start]
  54
  55 adjacents :: Position -> Slides -> Grid -> [Position]
  56 adjacents here slides forest = filter (`S.notMember` forest) $ fmap (here ^+^) deltas
  57   where deltas = case M.lookup here slides of
  58                   Nothing ->  [ V2 0 1, V2 1 0, V2 0 (-1), V2 (-1) 0 ]
  59                   Just SlideLeft -> [ V2 0 (-1) ]
  60                   Just SlideRight -> [ V2 0 1 ]
  61                   Just SlideUp -> [ V2 (-1) 0 ]
  62                   Just SlideDown -> [ V2 1 0 ]
  63
  64 searchStep :: Slides -> Grid -> [Position] -> [[Position]]
  65 searchStep _ _ [] = []
  66 searchStep slides forest path@(here:rest) = fmap (:path) valids
  67   where nexts = adjacents here slides forest
  68         valids = filter (`notElem` rest) nexts
  69
  70 search :: Slides -> Grid -> [Position] -> CompressedMap -> [[Position]] -> CompressedMap
  71 search _ _ _ foundPaths [] = foundPaths
  72 search slides forest goals foundPaths (current:agenda)
  73   | head current `elem` goals = search slides forest goals foundPaths' agenda
  74   | otherwise = search slides forest goals foundPaths (agenda ++ extendeds)
  75   where extendeds = searchStep slides forest current
  76         origin = last current
  77         foundPaths' = if origin == head current then foundPaths
  78                         else M.adjust (cp :) origin foundPaths
  79         cp = CPath (head current) (length current - 1)
  80
  81 -- collapsing the map
  82
  83 interestingPoints :: Slides -> Grid -> Position -> Position -> CompressedMap
  84 interestingPoints slides forest start end = M.fromList [(p, []) | p <- pointsSE]
  85   where Just minR = minimumOf (folded . _r) forest
  86         Just maxR = maximumOf (folded . _r) forest
  87         Just minC = minimumOf (folded . _c) forest
  88         Just maxC = maximumOf (folded . _c) forest
  89         points = [ V2 r c | r <- [(minR + 2)..(maxR - 2)]
  90                  , c <- [(minC + 1)..(maxC - 1)]
  91                  , (V2 r c) `S.notMember` forest
  92                  , (length $ adjacents (V2 r c) slides forest) > 2
  93                  ]
  94         pointsSE = start : end : points
  95
  96 compress :: Slides -> Grid -> Position -> Position -> CompressedMap
  97 compress slides forest start end = foldl' go compressed0 iPoints
  98   where compressed0 = interestingPoints slides forest start end
  99         iPoints = M.keys compressed0
 100         go com here = search slides forest iPoints com $ fmap (: [here]) $ adjacents here slides forest
 101
 102
 103 searchCompressed :: CompressedMap -> Position -> [Position] -> [[Position]]
 104 searchCompressed map goal current = runPar $ searchCompressedM parallelDepthLimit map goal current
 105
 106 searchCompressedM :: Int -> CompressedMap -> Position -> [Position] -> Par [[Position]]
 107 -- searchCompressed _ _ _ (c:_) _ | DT.trace (show c) False = undefined
 108 searchCompressedM depthLimit map goal current
 109   | head current == goal = return [current]
 110   | depthLimit == 0 = return $ searchCompressedTree map goal current
 111   | otherwise =
 112     do  paths <- parMapM (searchCompressedM (depthLimit - 1) map goal) nextPositions
 113         return $ concat paths
 114   where neighbours0 = map M.! (head current)
 115         neighbours = neighbours0 ^.. folded . filtered ((`notElem` current) . _nextPos)
 116         nextPositions = fmap ((: current) . _nextPos) neighbours
 117
 118 searchCompressedTree :: CompressedMap -> Position -> [Position] -> [[Position]]
 119 -- searchCompressed _ _ _ (c:_) _ | DT.trace (show c) False = undefined
 120 searchCompressedTree map goal current
 121   | head current == goal = [current]
 122   | otherwise = concatMap (searchCompressedTree map goal) nextPositions
 123   where neighbours0 = map M.! (head current)
 124         neighbours = neighbours0 ^.. folded . filtered ((`notElem` current) . _nextPos)
 125         nextPositions = fmap ((: current) . _nextPos) neighbours
 126
 127 pathLength :: CompressedMap -> [Position] -> Int
 128 pathLength map ps = sum $ zipWith (stepLength map) ps $ tail ps
 129
 130 stepLength :: CompressedMap -> Position -> Position -> Int
 131 stepLength map here there =
 132   -- head $ (map M.! there) ^.. folded . filtered ((== here) . _nextPos) . pathLen
 133   head $ (map M.! there) ^.. folded . filteredBy (nextPos . only here) . pathLen
 134
 135 -- reading the map
 136
 137 mkGrid :: String -> (Grid, Slides, Position, Position)
 138 mkGrid text = ((S.union forest caps), slides, start, end)
 139   where rows = lines text
 140         maxR = length rows - 1
 141         maxC = (length $ head rows) - 1
 142         forest = S.fromList [ V2 r c | r <- [0..maxR], c <- [0..maxC]
 143                             , rows !! r !! c == '#'
 144                             ]
 145         slides = M.fromList [ (V2 r c, readSlide (rows !! r !! c))
 146                             | r <- [0..maxR], c <- [0..maxC]
 147                             , elem (rows !! r !! c) ("<>^v" :: String)
 148                             ]
 149         start = head $ [ V2 0 c | c <- [0..maxC]
 150                        , rows !! 0 !! c == '.'
 151                        ]
 152         end = head $ [ V2 maxR c | c <- [0..maxC]
 153                      , rows !! maxR !! c == '.'
 154                      ]
 155         caps = S.fromList [start ^+^ (V2 -1 0), end ^+^ (V2 1 0)]
 156
 157 readSlide :: Char -> Slide
 158 readSlide '<' = SlideLeft
 159 readSlide '>' = SlideRight
 160 readSlide '^' = SlideUp
 161 readSlide 'v' = SlideDown