advent17/src/advent17.hs

   1 -- import Debug.Trace
   2
   3 import qualified Data.Set as S
   4 import Linear (V3(..), V4(..), (^+^))
   5
   6 class (Num a, Ord a) => Coord a where
   7     (^+^^) :: a -> a -> a
   8     neighbourCells :: S.Set a
   9 instance Coord (V3 Int) where
  10     x ^+^^ y  = x ^+^ y
  11     neighbourCells = S.fromList [ V3 dx dy dz
  12              | dx <- [-1, 0, 1]
  13              , dy <- [-1, 0, 1]
  14              , dz <- [-1, 0, 1]
  15              , (dx, dy, dz) /= (0, 0, 0)
  16              ]
  17 instance Coord (V4 Int) where
  18     x ^+^^ y  = x ^+^ y
  19     neighbourCells = S.fromList [ V4 dx dy dz dw
  20              | dx <- [-1, 0, 1]
  21              , dy <- [-1, 0, 1]
  22              , dz <- [-1, 0, 1]
  23              , dw <- [-1, 0, 1]
  24              , (dx, dy, dz, dw) /= (0, 0, 0, 0)
  25              ]
  26
  27 type Grid a = S.Set a
  28
  29 main :: IO ()
  30 main =
  31     do grid0 <- readGrid "data/advent17.txt"
  32        print $ part1 grid0
  33        print $ part2 grid0
  34
  35
  36 part1 grid0 = S.size finalGrid
  37   where finalGrid = head $ drop 6 $ iterate update grid0
  38
  39 part2 grid0 = S.size finalGrid
  40   where grid4 = conv34 grid0
  41         finalGrid = head $ drop 6 $ iterate update grid4
  42
  43
  44 readGrid :: String -> IO (Grid (V3 Int))
  45 readGrid filename =
  46     do  gs <- readFile filename
  47         let grid = lines gs
  48         let isActive x y = (grid!!y)!!x == '#'
  49         let maxX = length (head grid) - 1
  50         let maxY = length grid - 1
  51         return $ S.fromList [ V3 x y 0
  52                             | x <- [0..maxX], y <- [0..maxY], isActive x y]
  53
  54 conv34 :: Grid (V3 Int) -> Grid (V4 Int)
  55 conv34 grid = S.map conv34Cell grid
  56
  57 conv34Cell (V3 x y z) = V4 x y z 0
  58
  59
  60 neighbourSpaces :: Coord a => a -> Grid a
  61 neighbourSpaces here = S.map (here ^+^^) neighbourCells
  62
  63 countOccupiedNeighbours :: Coord a => a -> Grid a -> Int
  64 countOccupiedNeighbours cell grid =
  65     S.size $ S.intersection grid $ neighbourSpaces cell
  66
  67 cubeSurvives :: Coord a => Grid a -> a -> Bool
  68 cubeSurvives grid cell = alive && (nNbrs == 2 || nNbrs == 3)
  69   where alive = cell `S.member` grid
  70         nNbrs = countOccupiedNeighbours cell grid
  71
  72 cubeBorn :: Coord a => Grid a -> a -> Bool
  73 cubeBorn grid cell = dead && (nNbrs == 3)
  74   where dead = cell `S.notMember` grid
  75         nNbrs = countOccupiedNeighbours cell grid
  76
  77 update :: Coord a => Grid a -> Grid a
  78 update grid = S.union (S.filter (cubeSurvives grid) grid)
  79                       (S.filter (cubeBorn grid) empties)
  80   where empties = (S.foldr mergeEmpties S.empty grid) `S.difference` grid
  81         mergeEmpties cell acc = S.union acc $ neighbourSpaces cell