advent14/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/18/advent-of-code-2023-day-14/
   2
   3 import AoC
   4 import Data.List
   5 import qualified Data.Map.Strict as M
   6 import qualified Data.Array.Unboxed as U
   7 import qualified Data.Array.ST as A
   8 import Data.STRef
   9 import Control.Monad.ST
  10 import Control.Monad
  11 import Data.Ix
  12 import Linear (V2(..), (^+^))
  13 import qualified Data.Sequence as Q
  14 import Data.Sequence (Seq( (:|>), (:<|) ) )
  15
  16 type Position = V2 Int
  17 type Grid = U.UArray Position Bool
  18 type MGrid s = A.STUArray s Position Bool
  19 type Gaps s = STRef s (Q.Seq Position)
  20 type Cache = M.Map Grid Int
  21
  22 main :: IO ()
  23 main =
  24   do  dataFileName <- getDataFileName
  25       text <- readFile dataFileName
  26       let (rGrid, cGrid) = readGrids text
  27       print $ part1 rGrid cGrid
  28       print $ part2 rGrid cGrid
  29
  30
  31 part1, part2 :: Grid -> Grid -> Int
  32 part1 rGrid cGrid = getLoad $ rollNorth rGrid cGrid
  33 part2 rGrid cGrid = getLoad finalGrid
  34   where (grid', cache, repeatEnd) = findRepeat rGrid cGrid
  35         repeatStart = cache M.! grid'
  36         repeatLen = repeatEnd - repeatStart
  37         finalIndex = ((1e9 - repeatStart) `mod` repeatLen) + repeatStart
  38         (finalGrid, _) = M.findMin $ M.filter (== finalIndex) cache
  39
  40
  41 findRepeat :: Grid -> Grid -> (Grid, Cache, Int)
  42 findRepeat rGrid cGrid = head $ dropWhile test $ iterate go (rGrid, M.empty, 0)
  43   where test (g, c, _) = g `M.notMember` c
  44         go (g, c, i) = (rollCycle g cGrid, M.insert g i c, (i + 1))
  45
  46 rollNorth, rollCycle :: Grid -> Grid -> Grid
  47 rollNorth rGrid cGrid = roll [(V2 0 0, V2 0 1, V2 1 0)] cGrid rGrid
  48
  49 rollCycle rGrid cGrid = roll [ (V2 0 0, V2 0 1, V2 1 0)
  50                              , (V2 0 0, V2 1 0, V2 0 1)
  51                              , (V2 r 0, V2 0 1, V2 -1 0)
  52                              , (V2 0 c, V2 1 0, V2 0 -1)
  53                              ]
  54                              cGrid rGrid
  55   where (_, V2 r c) = U.bounds rGrid
  56
  57 roll :: [(Position, Position, Position)] -> Grid -> Grid -> Grid
  58 roll moveSpecs cGrid rGrid =
  59   A.runSTUArray $
  60     do grid <- A.thaw rGrid
  61        holes <- newSTRef Q.Empty
  62        forM_ moveSpecs $ \(start, majorStep, minorStep) ->
  63          forM_ (takeWhile (inBounds rGrid) $ iterate (^+^ majorStep) start) $ \maj ->
  64            do writeSTRef holes Q.Empty
  65               forM_ (takeWhile (inBounds rGrid) $ iterate (^+^ minorStep) maj) $ \here ->
  66                 rollPosition grid cGrid holes here
  67        return grid
  68
  69 rollPosition :: (MGrid s) -> Grid -> (Gaps s) -> Position -> ST s ()
  70 rollPosition grid cGrid holes here
  71   | cGrid U.! here = writeSTRef holes Q.Empty
  72   | otherwise = do roundHere <- A.readArray grid here
  73                    holesVal <- readSTRef holes
  74                    if roundHere then
  75                      case holesVal of
  76                        Q.Empty -> return ()
  77                        (h :<| hs) -> do A.writeArray grid h True
  78                                         A.writeArray grid here False
  79                                         writeSTRef holes (hs :|> here)
  80                    else modifySTRef holes (:|> here)
  81
  82 inBounds :: Grid -> Position -> Bool
  83 inBounds g h = inRange (U.bounds g) h
  84
  85 getLoad :: Grid -> Int
  86 getLoad grid = sum [getColLoad c | c <- [0..cMax]]
  87   where (_, V2 rMax cMax) = U.bounds grid
  88         getColLoad c = sum [(rMax - r + 1) | r <- [0..rMax], grid U.! (V2 r c)]
  89
  90 readGrids :: String -> (Grid, Grid)
  91 readGrids text = (rGrid, cGrid)
  92   where rows = lines text
  93         r = length rows - 1
  94         c = (length $ head rows) - 1
  95         rGrid = U.listArray ((V2 0 0), (V2 r c)) $ fmap (== 'O') $ concat rows
  96         cGrid = U.listArray ((V2 0 0), (V2 r c)) $ fmap (== '#') $ concat rows
  97
  98 showGrids :: Grid -> Grid -> String
  99 showGrids rGrid cGrid = unlines rows
 100   where (_, V2 rMax cMax) = U.bounds rGrid
 101         rows = [showRow r | r <- [0..rMax]]
 102         showRow r = [showElem r c | c <- [0..cMax]]
 103         showElem r c = let isR = rGrid U.! (V2 r c)
 104                            isC = cGrid U.! (V2 r c)
 105                        in if | isR && isC -> 'X'
 106                              | isR -> 'O'
 107                              | isC -> '#'
 108                              | otherwise -> '.'