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),
  40                                                   row = r, column = 5}
  41                                           | r <- [1..gridSize] ]
  42   | column == 1             = S.singleton ( Cell { level = (level - 1),
  43                                                     row = 3, column = 2})
  44   | otherwise               = S.singleton ( Cell { column = (column - 1), ..})
  45
  46 neighbourSpacesRight :: Cell -> Grid
  47 neighbourSpacesRight (Cell {..})
  48   | column == 2 && row == 3 = S.fromList  [ Cell { level = (level + 1),
  49                                                    row = r, column = 1}
  50                                           | r <- [1..gridSize] ]
  51   | column == 5             = S.singleton ( Cell { level = (level - 1),
  52                                                    row = 3, column = 4})
  53   | otherwise               = S.singleton ( Cell { column = (column + 1), ..})
  54
  55 neighbourSpacesAbove :: Cell -> Grid
  56 neighbourSpacesAbove (Cell {..})
  57   | row == 4 && column == 3 = S.fromList  [ Cell { level = (level + 1),
  58                                                    row = 5, column = c}
  59                                           | c <- [1..gridSize] ]
  60   | row == 1                = S.singleton ( Cell { level = (level - 1),
  61                                                    row = 2, column = 3})
  62   | otherwise               = S.singleton ( Cell { row = (row - 1), ..})
  63
  64 neighbourSpacesBelow :: Cell -> Grid
  65 neighbourSpacesBelow (Cell {..})
  66   | row == 2 && column == 3 = S.fromList  [ Cell { level = (level + 1),
  67                                                    row = 1, column = c}
  68                                           | c <- [1..gridSize] ]
  69   | row == 5                = S.singleton ( Cell { level = (level - 1),
  70                                                    row = 4, column = 3})
  71   | otherwise               = S.singleton ( Cell { row = (row + 1), ..})
  72
  73
  74 countOccupiedNeighbours :: Cell -> Grid -> Int
  75 countOccupiedNeighbours cell grid = S.size $ S.intersection grid $ neighbourSpaces cell
  76
  77 bugSurvives :: Grid -> Cell -> Bool
  78 bugSurvives grid cell = alive && oneNeighbour
  79   where alive = cell `S.member` grid
  80         oneNeighbour = (countOccupiedNeighbours cell grid) == 1
  81
  82 bugBorn :: Grid -> Cell -> Bool
  83 bugBorn grid cell = dead && (nNbrs == 1 || nNbrs == 2)
  84   where dead = cell `S.notMember` grid
  85         nNbrs = countOccupiedNeighbours cell grid
  86
  87 update :: Grid -> Grid
  88 update grid = S.union (S.filter (bugSurvives grid) bugs) (S.filter (bugBorn grid) empties)
  89   where bugs = grid
  90         empties = (S.foldr mergeEmpties S.empty grid) `S.difference` bugs
  91         mergeEmpties cell acc = S.union acc $ neighbourSpaces cell
  92