advent24/src/advent24b.hs

   1 import Debug.Trace
   2
   3 import qualified Data.Set as S
   4
   5 data Cell = Cell { level :: Int
   6                  , row :: Int
   7                  , column :: Int
   8                  } deriving (Show, Eq, Ord)
   9 type Grid = S.Set Cell
  10
  11
  12 gridSize = 5
  13
  14 main :: IO ()
  15 main =
  16     do grid0 <- readGrid
  17        print grid0
  18        let finalGrid = head $ drop 200 $ iterate update grid0
  19        print $ S.size finalGrid
  20
  21
  22 readGrid =
  23     do  gs <- readFile "data/advent24.txt"
  24         let grid = lines gs
  25         let isBug r c = (grid!!(r - 1))!!(c - 1) == '#'
  26         let level = 0
  27         return $ S.fromList [Cell {..} | row <- [1..gridSize], column <- [1..gridSize], isBug row column]
  28
  29 neighbourSpaces :: Cell -> Grid
  30 neighbourSpaces cell =
  31      (  (neighbourSpacesLeft cell)
  32      <> (neighbourSpacesRight cell)
  33      <> (neighbourSpacesAbove cell)
  34      <> (neighbourSpacesBelow cell)
  35      )
  36
  37 neighbourSpacesLeft :: Cell -> Grid
  38 neighbourSpacesLeft (Cell {..})
  39   | column == 4 && row == 3 = S.fromList  [ Cell { level = (level + 1), row = r, column = 5} | r <- [1..gridSize] ]
  40   | column == 1             = S.singleton ( Cell { level = (level - 1), row = 3, column = 2})
  41   | otherwise               = S.singleton ( Cell { level, row, column = (column - 1)})
  42
  43 neighbourSpacesRight :: Cell -> Grid
  44 neighbourSpacesRight (Cell {..})
  45   | column == 2 && row == 3 = S.fromList  [ Cell { level = (level + 1), row = r, column = 1} | r <- [1..gridSize] ]
  46   | column == 5             = S.singleton ( Cell { level = (level - 1), row = 3, column = 4})
  47   | otherwise               = S.singleton ( Cell { level, row, column = (column + 1)})
  48
  49 neighbourSpacesAbove :: Cell -> Grid
  50 neighbourSpacesAbove (Cell {..})
  51   | row == 4 && column == 3 = S.fromList  [ Cell { level = (level + 1), row = 5, column = c} | c <- [1..gridSize] ]
  52   | row == 1                = S.singleton ( Cell { level = (level - 1), row = 2, column = 3})
  53   | otherwise               = S.singleton ( Cell { level, row = (row - 1), column})
  54
  55 neighbourSpacesBelow :: Cell -> Grid
  56 neighbourSpacesBelow (Cell {..})
  57   | row == 2 && column == 3 = S.fromList  [ Cell { level = (level + 1), row = 1, column = c} | c <- [1..gridSize] ]
  58   | row == 5                = S.singleton ( Cell { level = (level - 1), row = 4, column = 3})
  59   | otherwise               = S.singleton ( Cell { level, row = (row + 1), column})
  60
  61
  62 countOccupiedNeighbours :: Cell -> Grid -> Int
  63 countOccupiedNeighbours cell grid = S.size $ S.intersection grid $ neighbourSpaces cell
  64
  65 bugSurvives :: Grid -> Cell -> Bool
  66 bugSurvives grid cell = alive && oneNeighbour
  67   where alive = cell `S.member` grid
  68         oneNeighbour = (countOccupiedNeighbours cell grid) == 1
  69
  70 bugBorn :: Grid -> Cell -> Bool
  71 bugBorn grid cell = dead && (nNbrs == 1 || nNbrs == 2)
  72   where dead = cell `S.notMember` grid
  73         nNbrs = countOccupiedNeighbours cell grid
  74
  75 update :: Grid -> Grid
  76 update grid = S.union (S.filter (bugSurvives grid) bugs) (S.filter (bugBorn grid) empties)
  77   where bugs = grid
  78         empties = (S.foldr mergeEmpties S.empty grid) `S.difference` bugs
  79         mergeEmpties cell acc = S.union acc $ neighbourSpaces cell
  80