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