advent17/src/advent17a.hs

   1 -- import Debug.Trace
   2
   3 import qualified Data.Set as S
   4 import Linear (V3(..), (^+^))
   5
   6 type Coord = V3 Int -- x, y, z
   7 type Grid = S.Set Coord
   8
   9 main :: IO ()
  10 main =
  11     do grid0 <- readGrid "data/advent17.txt"
  12        let finalGrid = head $ drop 6 $ iterate update grid0
  13        print $ S.size finalGrid
  14
  15
  16 readGrid :: String -> IO Grid
  17 readGrid filename =
  18     do  gs <- readFile filename
  19         let grid = lines gs
  20         let isActive x y = (grid!!y)!!x == '#'
  21         let maxX = length (head grid) - 1
  22         let maxY = length grid - 1
  23         return $ S.fromList [ V3 x y 0
  24                             | x <- [0..maxX], y <- [0..maxY], isActive x y]
  25
  26 neighbourSpaces :: Coord -> Grid
  27 neighbourSpaces here = S.map (here ^+^) nbrs
  28   where nbrs = S.fromList [ V3 dx dy dz
  29              | dx <- [-1, 0, 1]
  30              , dy <- [-1, 0, 1]
  31              , dz <- [-1, 0, 1]
  32              , (dx, dy, dz) /= (0, 0, 0)]
  33
  34 countOccupiedNeighbours :: Coord -> Grid -> Int
  35 countOccupiedNeighbours cell grid =
  36   S.size $ S.intersection grid $ neighbourSpaces cell
  37
  38 cubeSurvives :: Grid -> Coord -> Bool
  39 cubeSurvives grid cell = alive && (nNbrs == 2 || nNbrs == 3)
  40   where alive = cell `S.member` grid
  41         nNbrs = countOccupiedNeighbours cell grid
  42
  43 cubeBorn :: Grid -> Coord -> Bool
  44 cubeBorn grid cell = dead && (nNbrs == 3)
  45   where dead = cell `S.notMember` grid
  46         nNbrs = countOccupiedNeighbours cell grid
  47
  48 update :: Grid -> Grid
  49 update grid = S.union (S.filter (cubeSurvives grid) grid)
  50                       (S.filter (cubeBorn grid) empties)
  51   where empties = (S.foldr mergeEmpties S.empty grid) `S.difference` grid
  52         mergeEmpties cell acc = S.union acc $ neighbourSpaces cell