advent17/src/advent17.hs

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