1 -- Writeup at https://work.njae.me.uk/2021/12/09/advent-of-code-2021-day-9/
5 import Data.List (sort)
6 import qualified Data.Set as S
8 import Linear (V2(..), (^+^))
11 type Grid = Array Coord Int
12 type Basin = S.Set Coord
16 do text <- readFile "data/advent09a.txt"
17 let grid = mkGrid text
22 mkGrid :: String -> Grid
23 mkGrid text = listArray ((V2 0 0), (V2 r c)) $ map digitToInt $ concat rows
24 where rows = lines text
26 c = (length $ head rows) - 1
29 part1 grid = sum $ map (riskLevel grid) $ lowPoints grid
32 part2 grid = product $ take 3 ordSizes
33 where lows = lowPoints grid
34 sizes = map (basinSize grid) lows
35 ordSizes = reverse $ sort sizes
37 riskLevel :: Grid -> Coord -> Int
38 riskLevel grid here = grid ! here + 1
40 lowPoints :: Grid -> [Coord]
41 lowPoints grid = filter (isLow grid) $ indices grid
43 isLow :: Grid -> Coord -> Bool
44 isLow grid here = all (> this) nbrs
45 where nbrs = map (grid ! ) $ neighbours grid here
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
53 basinSize :: Grid -> Coord -> Int
54 basinSize grid basinSeed = S.size $ breadthFirstSearch grid (S.singleton basinSeed) S.empty
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
64 else S.insert here basin
65 agenda' = S.union candidates $ S.delete here agenda
68 neighbours :: Grid -> Coord -> [Coord]
69 neighbours grid here = filter (inRange (bounds grid))
71 | delta <- [V2 -1 0, V2 1 0, V2 0 -1, V2 0 1]