advent09/Main.hs

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