Faster day 14, with mutable unboxed arrays

diff --git a/advent-of-code23.cabal b/advent-of-code23.cabal
index a0b288042abc30c380fe18855a0aa5aa5128e368..5178f7dd2d89dfe27acef7a590015a5c3155e27d 100644 (file)
--- 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 1d0e71eb9313e5e7e3aaa1ab45911bcd90e467e3..d3b0c621feaeedc80e070b07b2237d36aa7fbe5b 100644 (file)
--- 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 (file)
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'