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 $ bounds grid
  19       print $ part1 grid
  20       print $ part2 grid
  21
  22 mkGrid :: String -> Grid
  23 mkGrid text = listArray ((V2 0 0), (V2 r c)) $ map digitToInt $ concat rows
  24   where rows = lines text
  25         r = length rows - 1
  26         c = (length $ head rows) - 1
  27
  28 part1 :: Grid -> Int
  29 part1 grid = sum $ map (riskLevel grid) $ lowPoints grid
  30
  31 part2 :: Grid -> Int
  32 part2 grid = product $ take 3 ordSizes
  33   where lows = lowPoints grid
  34         sizes = map (basinSize grid) lows
  35         ordSizes = reverse $ sort sizes
  36
  37 riskLevel :: Grid -> Coord -> Int
  38 riskLevel grid here = grid ! here + 1
  39
  40 lowPoints :: Grid -> [Coord]
  41 lowPoints grid = filter (isLow grid) $ indices grid
  42
  43 isLow :: Grid -> Coord -> Bool
  44 isLow grid here = all (> this) nbrs
  45   where nbrs = map (grid ! ) $ neighbours grid here
  46         this = grid ! here
  47
  48 higherNeighbours :: Grid -> Coord -> [Coord]
  49 higherNeighbours grid here = filter isHigher $ neighbours grid here
  50   where this = grid ! here
  51         isHigher there = (grid ! there) > this
  52
  53 basinSize :: Grid -> Coord -> Int
  54 basinSize grid basinSeed = S.size $ breadthFirstSearch grid (S.singleton basinSeed) S.empty
  55
  56 breadthFirstSearch :: Grid -> Basin -> Basin -> Basin
  57 breadthFirstSearch grid agenda basin
  58   | S.null agenda = basin
  59   | otherwise = breadthFirstSearch grid agenda' basin'
  60   where here = S.findMin agenda
  61         candidates = (S.fromList $ higherNeighbours grid here) \\ basin
  62         basin' = if (grid ! here) == 9
  63                  then basin
  64                  else S.insert here basin
  65         agenda' = S.union candidates $ S.delete here agenda
  66
  67
  68 neighbours :: Grid -> Coord -> [Coord]
  69 neighbours grid here = filter (inRange (bounds grid))
  70   [ here ^+^ delta
  71   | delta <- [V2 -1 0, V2 1 0, V2 0 -1, V2 0 1]
  72   ]
  73