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       -- putStrLn $ showGrid rGrid cGrid
  28       -- let rGrid' = rollNorth rGrid cGrid
  29       -- putStrLn $ showGrid rGrid' cGrid
  30       print $ part1 rGrid cGrid
  31       -- let rGrid1 = rollCycle rGrid cGrid
  32       -- putStrLn $ showGrid rGrid1 cGrid
  33       -- let rGrid2 = rollCycle rGrid1 cGrid
  34       -- putStrLn $ showGrid rGrid2 cGrid
  35       -- let rGrid3 = rollCycle rGrid2 cGrid
  36       -- putStrLn $ showGrid rGrid3 cGrid
  37       print $ part2 rGrid cGrid
  38
  39
  40 part1, part2 :: Grid -> Grid -> Int
  41 part1 rGrid cGrid = getLoad $ rollNorth rGrid cGrid
  42 part2 rGrid cGrid = getLoad finalGrid
  43   where (grid', cache, repeatEnd) = findRepeat rGrid cGrid
  44         repeatStart = cache M.! grid'
  45         repeatLen = repeatEnd - repeatStart
  46         finalIndex = ((1e9 - repeatStart) `mod` repeatLen) + repeatStart
  47         (finalGrid, _) = M.findMin $ M.filter (== finalIndex) cache
  48
  49
  50 findRepeat :: Grid -> Grid -> (Grid, Cache, Int)
  51 findRepeat rGrid cGrid = head $ dropWhile test $ iterate go (rGrid, M.empty, 0)
  52   where test (g, c, _) = g `M.notMember` c
  53         go (g, c, i) = (rollCycle g cGrid, M.insert g i c, (i + 1))
  54
  55 rollNorth, rollCycle :: Grid -> Grid -> Grid
  56 rollNorth rGrid cGrid = roll (V2 0 0) (V2 0 1) (V2 1 0) cGrid rGrid
  57
  58 rollCycle rGrid cGrid = foldl' go rGrid [ (V2 0 0, V2 0 1, V2 1 0)
  59                                         , (V2 0 0, V2 1 0, V2 0 1)
  60                                         , (V2 r 0, V2 0 1, V2 -1 0)
  61                                         , (V2 0 c, V2 1 0, V2 0 -1)
  62                                         ]
  63   where (_, V2 r c) = U.bounds rGrid
  64         go g (start, majorStep, minorStep) =
  65           roll start majorStep minorStep cGrid g
  66
  67 roll :: Position -> Position -> Position -> Grid -> Grid -> Grid
  68 roll start majorStep minorStep cGrid rGrid =
  69   A.runSTUArray $
  70     do grid <- A.thaw rGrid
  71        holes <- newSTRef Q.Empty
  72        forM_ (takeWhile (inBounds rGrid) $ iterate (^+^ majorStep) start) $ \maj ->
  73          do writeSTRef holes Q.Empty
  74             forM_ (takeWhile (inBounds rGrid) $ iterate (^+^ minorStep) maj) $ \here ->
  75               rollPosition grid cGrid holes here
  76        return grid
  77
  78 rollPosition :: (MGrid s) -> Grid -> (Gaps s) -> Position -> ST s ()
  79 rollPosition grid cGrid holes here
  80   | cGrid U.! here = writeSTRef holes Q.Empty
  81   | otherwise = do roundHere <- A.readArray grid here
  82                    holesVal <- readSTRef holes
  83                    if roundHere then
  84                      case holesVal of
  85                        Q.Empty -> return ()
  86                        (h :<| hs) -> do A.writeArray grid h True
  87                                         A.writeArray grid here False
  88                                         writeSTRef holes (hs :|> here)
  89                    else modifySTRef holes (:|> here)
  90
  91 inBounds :: Grid -> Position -> Bool
  92 inBounds g h = inRange (U.bounds g) h
  93
  94 getLoad :: Grid -> Int
  95 getLoad grid = sum columnLoads
  96   where (_, V2 rMax cMax) = U.bounds grid
  97         columnLoads = [getColLoad c | c <- [0..cMax]]
  98         getColLoad c = sum [(rMax - r + 1) | r <- [0..rMax], grid U.! (V2 r c)]
  99
 100 readGrids :: String -> (Grid, Grid)
 101 readGrids text = (rGrid, cGrid)
 102   where rows = lines text
 103         r = length rows - 1
 104         c = (length $ head rows) - 1
 105         rGrid = U.listArray ((V2 0 0), (V2 r c)) $ fmap (== 'O') $ concat rows
 106         cGrid = U.listArray ((V2 0 0), (V2 r c)) $ fmap (== '#') $ concat rows
 107
 108 -- readElem :: Char -> Element
 109 -- readElem '.' = Empty
 110 -- readElem '#' = Cube
 111 -- readElem 'O' = Round
 112
 113 showGrid :: Grid -> Grid -> String
 114 showGrid rGrid cGrid = unlines rows
 115   where (_, V2 rMax cMax) = U.bounds rGrid
 116         rows = [showRow r | r <- [0..rMax]]
 117         showRow r = [showElem r c | c <- [0..cMax]]
 118         showElem r c = let isR = rGrid U.! (V2 r c)
 119                            isC = cGrid U.! (V2 r c)
 120                        in if | isR && isC -> 'X'
 121                              | isR -> 'O'
 122                              | isC -> '#'
 123                              | otherwise -> '.'