1 -- Writeup at https://work.njae.me.uk/2021/12/09/advent-of-code-2021-day-8/
7 import Data.List hiding ((\\))
8 import qualified Data.Set as S
12 type Coord = (Int, Int) -- (row, column)
13 type Grid = Array Coord Int
14 type Basin = S.Set Coord
18 do text <- readFile "data/advent09.txt"
19 let grid = mkGrid text
22 -- print $ lowPoints grid
23 -- print $ breadthFirstSearch grid (S.singleton (4, 6)) S.empty
26 mkGrid :: String -> Grid
27 mkGrid text = listArray ((0, 0), (r, c)) $ map digitToInt $ concat rows
28 where rows = lines text
30 c = (length $ head rows) - 1
33 part1 grid = sum $ map (riskLevel grid) $ lowPoints grid
36 part2 grid = product $ take 3 ordSizes
37 where lows = lowPoints grid
38 sizes = map (basinSize grid) lows
39 ordSizes = reverse $ sort sizes
41 riskLevel :: Grid -> Coord -> Int
42 riskLevel grid here = grid ! here + 1
44 lowPoints :: Grid -> [Coord]
45 lowPoints grid = filter (isLow grid) $ indices grid
47 isLow :: Grid -> Coord -> Bool
48 isLow grid here = all (> this) nbrs
49 where nbrs = map (grid ! ) $ neighbours grid here
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
57 basinSize :: Grid -> Coord -> Int
58 basinSize grid basinSeed = S.size $ breadthFirstSearch grid (S.singleton basinSeed) S.empty
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
68 else S.insert here basin
69 agenda' = S.union candidates $ S.delete here agenda
72 neighbours :: Grid -> Coord -> [Coord]
73 neighbours grid (r, c) = filter (gValid grid)
75 | (dr, dc) <- [(-1, 0), (1, 0), (0, -1), (0, 1)]
78 gValid :: Grid -> Coord -> Bool
85 where ((minR, minC), (maxR, maxC)) = bounds grid