/ 2021 / 11 / p1.hs
p1.hs
 1  import Control.Monad
 2  import Control.Monad.ST
 3  import Data.Array.Unboxed
 4  import Data.Array.ST
 5  import Data.STRef
 6  import Data.Char
 7  import Data.Set qualified as Set
 8  
 9  readGrid :: String -> Array (Int, Int) Int
10  readGrid contents = listArray ((1, 1), (height, width)) $ concat input
11    where
12      input = fmap (fmap digitToInt) $ lines contents
13      width = length $ input !! 0
14      height = length input
15  
16  adjacent (y, x) = [(y + dy, x + dx) | dx <- [-1, 0, 1], dy <- [-1, 0, 1], (dx, dy) /= (0, 0)]
17  
18  thaw_ :: Array (Int, Int) Int -> ST s (STUArray s (Int, Int) Int)
19  thaw_ = thaw
20  
21  answer input = runST $ do
22    grid <- thaw_ input
23    flashes <- newSTRef 0
24    forM_ [1..100] $ \_ -> do
25      flashed <- consume flashes Set.empty (indices input) grid
26      forM_ (indices input) $ \i -> modifyArray' grid i (\v -> if v >= 10 then 0 else v)
27    readSTRef flashes
28    where
29      valid = inRange (bounds input)
30  
31      consume count flashed [] _ = return flashed
32      consume count flashed (p:ps) grid
33        | Set.member p flashed = consume count flashed ps grid
34        | otherwise = do
35          here <- readArray grid p
36          writeArray grid p (here + 1)
37          if here + 1 == 10
38            then do
39              modifySTRef' count (+ 1)
40              consume count (Set.insert p flashed) (filter valid (adjacent p) ++ ps) grid
41            else consume count flashed ps grid
42  
43  main = getContents >>= print . answer . readGrid