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 Data.Semigroup
   6 import Data.Monoid
   7 import qualified Data.Map.Strict as M
   8
   9 data Element = Empty | Cube | Round deriving (Show, Eq, Ord)
  10 type Grid = [[Element]]
  11
  12 type Cache = M.Map Grid Int
  13
  14
  15 main :: IO ()
  16 main =
  17   do  dataFileName <- getDataFileName
  18       text <- readFile dataFileName
  19       let grid = transpose $ fmap (fmap readElem) $ lines text
  20       -- print $ showGrid grid
  21       print $ part1 grid
  22       print $ part2 grid
  23
  24 part1, part2 :: Grid -> Int
  25 part1 grid = scoreGrid grid'
  26   where grid' = rollToCompletion grid
  27
  28 part2 grid = scoreGrid finalGrid
  29   where (grid', cache, repeatEnd) = findRepeat grid
  30         repeatStart = cache M.! grid'
  31         repeatLen = repeatEnd - repeatStart
  32         finalIndex = ((1e9 - repeatStart) `mod` repeatLen) + repeatStart
  33         (finalGrid, _) = M.findMin $ M.filter (== finalIndex) cache
  34
  35 readElem :: Char -> Element
  36 readElem '.' = Empty
  37 readElem '#' = Cube
  38 readElem 'O' = Round
  39
  40 rollToCompletion :: Grid -> Grid
  41 rollToCompletion grid = fst $ head $ dropWhile (uncurry (/=)) $ zip states $ tail states
  42   where states = iterate rollGrid grid
  43
  44 rollGrid :: Grid -> Grid
  45 rollGrid = fmap roll
  46
  47 roll :: [Element] -> [Element]
  48 roll [] = []
  49 roll (l:ls) = rs ++ [r]
  50   where (rs, r) = foldl' rollStep ([], l) ls
  51
  52 rollStep :: ([Element], Element) -> Element -> ([Element], Element)
  53 rollStep (handled, Empty) Round = (handled ++ [Round], Empty)
  54 rollStep (handled, target) source = (handled ++ [target], source)
  55
  56 scoreGrid :: Grid -> Int
  57 scoreGrid grid = sum $ fmap scoreRow indexedGrid
  58   where indexedGrid = zip [1..] $ reverse $ transpose grid
  59         scoreRow (i, r) = i * (length $ filter (== Round) r)
  60
  61 rotate1 :: Grid -> Grid
  62 rotate1 = transpose . fmap reverse
  63
  64 rollCycle :: Grid -> Grid
  65 rollCycle = appEndo (stimes 4 (Endo rotate1 <> Endo rollToCompletion))
  66
  67 findRepeat :: Grid -> (Grid, Cache, Int)
  68 findRepeat grid = head $ dropWhile test $ iterate go (grid, M.empty, 0)
  69   where test (g, c, _) = g `M.notMember` c
  70         go (g, c, i) = (rollCycle g, M.insert g i c, (i + 1))
  71
  72 showGrid :: Grid -> String
  73 showGrid grid = unlines $ fmap (fmap showElem) $ transpose grid
  74   where showElem Empty = '.'
  75         showElem Cube = '#'
  76         showElem Round = 'O'