From ac1cb7b87194225d518bae8e5a9d9efe7d27220a Mon Sep 17 00:00:00 2001 From: Neil Smith Date: Mon, 22 Jan 2024 15:14:59 +0000 Subject: [PATCH] Faster day 14, with mutable unboxed arrays --- advent-of-code23.cabal | 6 +- advent14/Main.hs | 165 +++++++++++++++++++++++++++-------------- advent14/MainSlow.hs | 76 +++++++++++++++++++ 3 files changed, 191 insertions(+), 56 deletions(-) create mode 100644 advent14/MainSlow.hs diff --git a/advent-of-code23.cabal b/advent-of-code23.cabal index a0b2880..5178f7d 100644 --- a/advent-of-code23.cabal +++ b/advent-of-code23.cabal @@ -185,8 +185,12 @@ executable advent13 executable advent14 import: common-extensions, build-directives main-is: advent14/Main.hs + build-depends: containers, array, linear +executable advent14slow + import: common-extensions, build-directives + main-is: advent14/MainSlow.hs build-depends: containers - + executable advent15 import: common-extensions, build-directives main-is: advent15/Main.hs diff --git a/advent14/Main.hs b/advent14/Main.hs index 1d0e71e..d3b0c62 100644 --- a/advent14/Main.hs +++ b/advent14/Main.hs @@ -2,75 +2,130 @@ import AoC import Data.List -import Data.Semigroup -import Data.Monoid -import qualified Data.Map.Strict as M +import qualified Data.Map.Strict as Map +import qualified Data.Array.Unboxed as U +import qualified Data.Array.ST as M +import Data.STRef +import Control.Monad.ST +import Control.Monad +import Data.Ix + + +import Linear (V2(..), (^+^)) +import qualified Data.Sequence as Q +import Data.Sequence (Seq( (:|>), (:<|) ) ) data Element = Empty | Cube | Round deriving (Show, Eq, Ord) -type Grid = [[Element]] -type Cache = M.Map Grid Int +type Position = V2 Int +type Grid = U.UArray Position Bool +type MGrid s = M.STUArray s Position Bool + +type Gaps s = STRef s (Q.Seq Position) + +type Cache = Map.Map Grid Int main :: IO () main = do dataFileName <- getDataFileName text <- readFile dataFileName - let grid = transpose $ fmap (fmap readElem) $ lines text - -- print $ showGrid grid - print $ part1 grid - print $ part2 grid - -part1, part2 :: Grid -> Int -part1 grid = scoreGrid grid' - where grid' = rollToCompletion grid - -part2 grid = scoreGrid finalGrid - where (grid', cache, repeatEnd) = findRepeat grid - repeatStart = cache M.! grid' + let (rGrid, cGrid) = readGrids text + -- putStrLn $ showGrid rGrid cGrid + -- let rGrid' = rollNorth rGrid cGrid + -- putStrLn $ showGrid rGrid' cGrid + print $ part1 rGrid cGrid + -- let rGrid1 = rollCycle rGrid cGrid + -- putStrLn $ showGrid rGrid1 cGrid + -- let rGrid2 = rollCycle rGrid1 cGrid + -- putStrLn $ showGrid rGrid2 cGrid + -- let rGrid3 = rollCycle rGrid2 cGrid + -- putStrLn $ showGrid rGrid3 cGrid + print $ part2 rGrid cGrid + + +part1 :: Grid -> Grid -> Int +part1 rGrid cGrid = getLoad $ rollNorth rGrid cGrid + +part2 rGrid cGrid = getLoad finalGrid + where (grid', cache, repeatEnd) = findRepeat rGrid cGrid + repeatStart = cache Map.! grid' repeatLen = repeatEnd - repeatStart finalIndex = ((1e9 - repeatStart) `mod` repeatLen) + repeatStart - (finalGrid, _) = M.findMin $ M.filter (== finalIndex) cache + (finalGrid, _) = Map.findMin $ Map.filter (== finalIndex) cache + + +findRepeat :: Grid -> Grid -> (Grid, Cache, Int) +findRepeat rGrid cGrid = head $ dropWhile test $ iterate go (rGrid, Map.empty, 0) + where test (g, c, _) = g `Map.notMember` c + go (g, c, i) = (rollCycle g cGrid, Map.insert g i c, (i + 1)) + +rollNorth, rollCycle :: Grid -> Grid -> Grid +rollNorth rGrid cGrid = roll (V2 0 0) (V2 0 1) (V2 1 0) cGrid rGrid + +rollCycle rGrid cGrid = foldl' go rGrid [ (V2 0 0, V2 0 1, V2 1 0) + , (V2 0 0, V2 1 0, V2 0 1) + , (V2 r 0, V2 0 1, V2 -1 0) + , (V2 0 c, V2 1 0, V2 0 -1) + ] + where (_, V2 r c) = U.bounds rGrid + go g (start, majorStep, minorStep) = + roll start majorStep minorStep cGrid g + +roll :: Position -> Position -> Position -> Grid -> Grid -> Grid +roll start majorStep minorStep cGrid rGrid = + M.runSTUArray $ + do grid <- M.thaw rGrid + holes <- newSTRef Q.Empty + forM_ (takeWhile (inBounds rGrid) $ iterate (^+^ majorStep) start) $ \maj -> + do writeSTRef holes Q.Empty + forM_ (takeWhile (inBounds rGrid) $ iterate (^+^ minorStep) maj) $ \here -> + do rollPosition grid cGrid holes here + return grid + +rollPosition :: (MGrid s) -> Grid -> (Gaps s) -> Position -> ST s () +rollPosition grid cGrid holes here + | cGrid U.! here = writeSTRef holes Q.Empty + | otherwise = do roundHere <- M.readArray grid here + holesVal <- readSTRef holes + if roundHere then + case holesVal of + Q.Empty -> return () + (h :<| hs) -> do M.writeArray grid h True + M.writeArray grid here False + writeSTRef holes (hs :|> here) + else modifySTRef holes (:|> here) + +inBounds :: Grid -> Position -> Bool +inBounds g h = inRange (U.bounds g) h + +getLoad :: Grid -> Int +getLoad grid = sum columnLoads + where (_, V2 rMax cMax) = U.bounds grid + columnLoads = [getColLoad c | c <- [0..cMax]] + getColLoad c = sum [(rMax - r + 1) | r <- [0..rMax], grid U.! (V2 r c)] + +readGrids :: String -> (Grid, Grid) +readGrids text = (rGrid, cGrid) + where rows = lines text + r = length rows - 1 + c = (length $ head rows) - 1 + rGrid = U.listArray ((V2 0 0), (V2 r c)) $ map ((== Round) . readElem) $ concat rows + cGrid = U.listArray ((V2 0 0), (V2 r c)) $ map ((== Cube ) . readElem) $ concat rows readElem :: Char -> Element readElem '.' = Empty readElem '#' = Cube readElem 'O' = Round -rollToCompletion :: Grid -> Grid -rollToCompletion grid = fst $ head $ dropWhile (uncurry (/=)) $ zip states $ tail states - where states = iterate rollGrid grid - -rollGrid :: Grid -> Grid -rollGrid = fmap roll - -roll :: [Element] -> [Element] -roll [] = [] -roll (l:ls) = rs ++ [r] - where (rs, r) = foldl' rollStep ([], l) ls - -rollStep :: ([Element], Element) -> Element -> ([Element], Element) -rollStep (handled, Empty) Round = (handled ++ [Round], Empty) -rollStep (handled, target) source = (handled ++ [target], source) - -scoreGrid :: Grid -> Int -scoreGrid grid = sum $ fmap scoreRow indexedGrid - where indexedGrid = zip [1..] $ reverse $ transpose grid - scoreRow (i, r) = i * (length $ filter (== Round) r) - -rotate1 :: Grid -> Grid -rotate1 = transpose . fmap reverse - -rollCycle :: Grid -> Grid -rollCycle = appEndo (stimes 4 (Endo rotate1 <> Endo rollToCompletion)) - -findRepeat :: Grid -> (Grid, Cache, Int) -findRepeat grid = head $ dropWhile test $ iterate go (grid, M.empty, 0) - where test (g, c, _) = g `M.notMember` c - go (g, c, i) = (rollCycle g, M.insert g i c, (i + 1)) - -showGrid :: Grid -> String -showGrid grid = unlines $ fmap (fmap showElem) $ transpose grid - where showElem Empty = '.' - showElem Cube = '#' - showElem Round = 'O' +showGrid :: Grid -> Grid -> String +showGrid rGrid cGrid = unlines rows + where (_, V2 rMax cMax) = U.bounds rGrid + rows = [showRow r | r <- [0..rMax]] + showRow r = [showElem r c | c <- [0..cMax]] + showElem r c = let isR = rGrid U.! (V2 r c) + isC = cGrid U.! (V2 r c) + in if | isR && isC -> 'X' + | isR -> 'O' + | isC -> '#' + | otherwise -> '.' diff --git a/advent14/MainSlow.hs b/advent14/MainSlow.hs new file mode 100644 index 0000000..1d0e71e --- /dev/null +++ b/advent14/MainSlow.hs @@ -0,0 +1,76 @@ +-- Writeup at https://work.njae.me.uk/2023/12/18/advent-of-code-2023-day-14/ + +import AoC +import Data.List +import Data.Semigroup +import Data.Monoid +import qualified Data.Map.Strict as M + +data Element = Empty | Cube | Round deriving (Show, Eq, Ord) +type Grid = [[Element]] + +type Cache = M.Map Grid Int + + +main :: IO () +main = + do dataFileName <- getDataFileName + text <- readFile dataFileName + let grid = transpose $ fmap (fmap readElem) $ lines text + -- print $ showGrid grid + print $ part1 grid + print $ part2 grid + +part1, part2 :: Grid -> Int +part1 grid = scoreGrid grid' + where grid' = rollToCompletion grid + +part2 grid = scoreGrid finalGrid + where (grid', cache, repeatEnd) = findRepeat grid + repeatStart = cache M.! grid' + repeatLen = repeatEnd - repeatStart + finalIndex = ((1e9 - repeatStart) `mod` repeatLen) + repeatStart + (finalGrid, _) = M.findMin $ M.filter (== finalIndex) cache + +readElem :: Char -> Element +readElem '.' = Empty +readElem '#' = Cube +readElem 'O' = Round + +rollToCompletion :: Grid -> Grid +rollToCompletion grid = fst $ head $ dropWhile (uncurry (/=)) $ zip states $ tail states + where states = iterate rollGrid grid + +rollGrid :: Grid -> Grid +rollGrid = fmap roll + +roll :: [Element] -> [Element] +roll [] = [] +roll (l:ls) = rs ++ [r] + where (rs, r) = foldl' rollStep ([], l) ls + +rollStep :: ([Element], Element) -> Element -> ([Element], Element) +rollStep (handled, Empty) Round = (handled ++ [Round], Empty) +rollStep (handled, target) source = (handled ++ [target], source) + +scoreGrid :: Grid -> Int +scoreGrid grid = sum $ fmap scoreRow indexedGrid + where indexedGrid = zip [1..] $ reverse $ transpose grid + scoreRow (i, r) = i * (length $ filter (== Round) r) + +rotate1 :: Grid -> Grid +rotate1 = transpose . fmap reverse + +rollCycle :: Grid -> Grid +rollCycle = appEndo (stimes 4 (Endo rotate1 <> Endo rollToCompletion)) + +findRepeat :: Grid -> (Grid, Cache, Int) +findRepeat grid = head $ dropWhile test $ iterate go (grid, M.empty, 0) + where test (g, c, _) = g `M.notMember` c + go (g, c, i) = (rollCycle g, M.insert g i c, (i + 1)) + +showGrid :: Grid -> String +showGrid grid = unlines $ fmap (fmap showElem) $ transpose grid + where showElem Empty = '.' + showElem Cube = '#' + showElem Round = 'O' -- 2.34.1