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