advent09/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/09/advent-of-code-2021-day-8/
   2
   3
   4 import Data.Array
   5 import Data.Char
   6 -- import Data.Maybe
   7 import Data.List hiding ((\\))
   8 import qualified Data.Set as S
   9 import Data.Set ((\\))
  10
  11
  12 type Coord = (Int, Int) -- (row, column)
  13 type Grid = Array Coord Int
  14 type Basin = S.Set Coord
  15
  16 main :: IO ()
  17 main =
  18   do  text <- readFile "data/advent09.txt"
  19       let grid = mkGrid text
  20       print $ bounds grid
  21       print $ part1 grid
  22       -- print $ lowPoints grid
  23       -- print $ breadthFirstSearch grid (S.singleton (4, 6)) S.empty
  24       print $ part2 grid
  25
  26 mkGrid :: String -> Grid
  27 mkGrid text = listArray ((0, 0), (r, c)) $ map digitToInt $ concat rows
  28   where rows = lines text
  29         r = length rows - 1
  30         c = (length $ head rows) - 1
  31
  32 part1 :: Grid -> Int
  33 part1 grid = sum $ map (riskLevel grid) $ lowPoints grid
  34
  35 part2 :: Grid -> Int
  36 part2 grid = product $ take 3 ordSizes
  37   where lows = lowPoints grid
  38         sizes = map (basinSize grid) lows
  39         ordSizes = reverse $ sort sizes
  40
  41 riskLevel :: Grid -> Coord -> Int
  42 riskLevel grid here = grid ! here + 1
  43
  44 lowPoints :: Grid -> [Coord]
  45 lowPoints grid = filter (isLow grid) $ indices grid
  46
  47 isLow :: Grid -> Coord -> Bool
  48 isLow grid here = all (> this) nbrs
  49   where nbrs = map (grid ! ) $ neighbours grid here
  50         this = grid ! here
  51
  52 higherNeighbours :: Grid -> Coord -> [Coord]
  53 higherNeighbours grid here = filter isHigher $ neighbours grid here
  54   where this = grid ! here
  55         isHigher there = (grid ! there) > this
  56
  57 basinSize :: Grid -> Coord -> Int
  58 basinSize grid basinSeed = S.size $ breadthFirstSearch grid (S.singleton basinSeed) S.empty
  59
  60 breadthFirstSearch :: Grid -> Basin -> Basin -> Basin
  61 breadthFirstSearch grid agenda basin
  62   | S.null agenda = basin
  63   | otherwise = breadthFirstSearch grid agenda' basin'
  64   where here = S.findMin agenda
  65         candidates = (S.fromList $ higherNeighbours grid here) \\ basin
  66         basin' = if (grid ! here) == 9
  67                  then basin
  68                  else S.insert here basin
  69         agenda' = S.union candidates $ S.delete here agenda
  70
  71
  72 neighbours :: Grid -> Coord -> [Coord]
  73 neighbours grid (r, c) = filter (gValid grid)
  74   [ (r + dr, c + dc)
  75   | (dr, dc) <- [(-1, 0), (1, 0), (0, -1), (0, 1)]
  76   ]
  77
  78 gValid :: Grid -> Coord -> Bool
  79 gValid grid (r, c)
  80   | r < minR = False
  81   | c < minC = False
  82   | r > maxR = False
  83   | c > maxC = False
  84   | otherwise = True
  85   where ((minR, minC), (maxR, maxC)) = bounds grid