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    flashAll 1 grid
24    where
25      total = rangeSize (bounds input)
26      valid = inRange (bounds input)
27  
28      flashAll i grid = do
29        flashes <- step grid
30        if Set.size flashes == total then return i else flashAll (i + 1) grid
31  
32      step grid = do
33        flashed <- consume Set.empty (indices input) grid
34        forM_ (indices input) $ \i -> modifyArray' grid i (\v -> if v >= 10 then 0 else v)
35        return flashed
36  
37      consume flashed [] _ = return flashed
38      consume flashed (p:ps) grid
39        | Set.member p flashed = consume flashed ps grid
40        | otherwise = do
41          here <- readArray grid p
42          writeArray grid p (here + 1)
43          if here + 1 == 10
44            then consume (Set.insert p flashed) (filter valid (adjacent p) ++ ps) grid
45            else consume flashed ps grid
46  
47  main = getContents >>= print . answer . readGrid