advent11/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/12/advent-of-code-2021-day-11/
   2
   3 import Data.Array.IArray
   4 import Data.Char
   5 import Linear (V2(..), (^+^))
   6
   7 type Coord = V2 Int
   8 type Grid = Array Coord Octopus
   9
  10 data Octopus = Octopus Int Bool
  11   deriving (Ord, Eq, Show)
  12
  13 main :: IO ()
  14 main =
  15   do  text <- readFile "data/advent11.txt"
  16       let grid = mkGrid text
  17       print $ part1 grid
  18       print $ part2 grid
  19
  20 mkGrid :: String -> Grid
  21 mkGrid text = listArray ((V2 0 0), (V2 r c)) $ map mkOct $ concat rows
  22   where rows = lines text
  23         r = length rows - 1
  24         c = (length $ head rows) - 1
  25         mkOct e = Octopus (digitToInt e) False
  26
  27 part1 :: Grid -> Int
  28 part1 grid = snd $ (simulate grid) !! 100
  29
  30 part2 :: Grid -> Int
  31 part2 grid = length $ takeWhile notSyncronised $ simulate grid
  32   where notSyncronised (g, _) = not $ simultaneous g
  33
  34 simulate :: Grid -> [(Grid, Int)]
  35 simulate grid = iterate step (grid, 0)
  36
  37 step :: (Grid, Int) -> (Grid, Int)
  38 step (grid0, flashCount0) = (grid3, flashCount0 + numFlashers)
  39   where grid1 = increment grid0
  40         triggers = findFlashers grid1
  41         grid2 = flash grid1 triggers
  42         flashers = findFlashers grid2
  43         numFlashers = length flashers
  44         grid3 = resetFlashers grid2 flashers
  45
  46 simultaneous :: Grid -> Bool
  47 simultaneous grid = all zeroOct $ elems grid
  48   where zeroOct (Octopus 0 _) = True
  49         zeroOct _ = False
  50
  51 increment :: Grid -> Grid
  52 increment = amap incrementOne
  53
  54 incrementSome :: Grid -> [Coord] -> Grid
  55 incrementSome grid locations = grid // (zip locations incrementedOcts)
  56   where incrementedOcts = map (incrementOne . (grid !)) locations
  57
  58 incrementOne :: Octopus -> Octopus
  59 incrementOne (Octopus energy flashed) = Octopus (energy + 1) flashed
  60
  61 findFlashers :: Grid -> [Coord]
  62 findFlashers = map fst . filter (overpowered . snd) . assocs
  63   where overpowered (Octopus energy _) = energy > 9
  64
  65 flash :: Grid -> [Coord] -> Grid
  66 flash grid [] = grid
  67 flash grid (here:agenda)
  68   -- already flashed, so ignore
  69   | flashed == True = flash grid agenda
  70   -- not enough energy to flash, so ignore
  71   | energy <= 9 = flash grid agenda
  72   | otherwise = flash grid'' agenda'
  73   where Octopus energy flashed = grid ! here
  74         nbrs = neighbours grid here
  75         -- set this as flashed
  76         octopus' = Octopus energy True
  77         grid' = grid // [(here, octopus')]
  78         -- add negighbours to agenda
  79         agenda' = nbrs ++ agenda
  80         -- increment neighbours
  81         grid'' = incrementSome grid' nbrs
  82
  83 resetFlashers :: Grid -> [Coord] -> Grid
  84 resetFlashers grid locations = grid // (zip locations resetOcts)
  85   where resetOcts = repeat (Octopus 0 False)
  86
  87 neighbours :: Grid -> Coord -> [Coord]
  88 neighbours grid here = filter (inRange (bounds grid))
  89   [ here ^+^ (V2 r c)
  90   | r <- [-1, 0, 1]
  91   , c <- [-1, 0, 1]
  92   , (r, c) /= (0, 0)
  93   ]