The hardest part of this puzzle was deciding on the best data structure and library for the grid.
Data types, the grid, and the Ix data class
Haskell has a surfeit of array types. There's the standard Array library, the Vector library that can be misused to provide multi-dimensional arrays, and more complex libraries like Repa, Massiv, and Dense. There's even the Grid library for representing grids, but that only holds locations, not data. The problem with Repa, Massiv, and Dense is that the have very complex types to account for controlling how and when evaluation takes place, complicated indexing schemes, and similar. That was all overkill for this problem, even if it would have been fun to learn about.
In the end, I settled on the basic Array library, but using the
V2 data type from Linear for the coordinates.
type Coord = V2 Int type Grid = Array Coord Int type Basin = S.Set Coord
Making the grid from the input was fairly simple.
mkGrid :: String -> Grid mkGrid text = listArray ((V2 0 0), (V2 r c)) $ map digitToInt $ concat rows where rows = lines text r = length rows - 1 c = (length $ head rows) - 1
Finding neighbours was made much simpler by a little investigation into the
Ix data class, and the use of
inRange functions to tell that a possible location is within the array.
neighbours :: Grid -> Coord -> [Coord] neighbours grid here = filter (inRange (bounds grid)) [ here ^+^ delta | delta <- [V2 -1 0, V2 1 0, V2 0 -1, V2 0 1] ]
From that, finding the the lowest points is examining each location in the grid and checking that all its neighbours are higher.
part1 :: Grid -> Int part1 grid = sum $ map (riskLevel grid) $ lowPoints grid riskLevel :: Grid -> Coord -> Int riskLevel grid here = grid ! here + 1 lowPoints :: Grid -> [Coord] lowPoints grid = filter (isLow grid) $ indices grid isLow :: Grid -> Coord -> Bool isLow grid here = all (> this) nbrs where nbrs = map (grid ! ) $ neighbours grid here this = grid ! here
This is a by-now standard flood fill / breadth-first search algorithm. Starting from a lowest point, it adds all higher points to the boundary / agenda (if they're not already in the basin), adds the current point to the basin, and repeats. There's a bit of fiddling around with not adding positions at height 9, and I don't think it would work at tracking across plateaus of height below 9. But it gives the right answer.
part2 :: Grid -> Int part2 grid = product $ take 3 ordSizes where lows = lowPoints grid sizes = map (basinSize grid) lows ordSizes = reverse $ sort sizes higherNeighbours :: Grid -> Coord -> [Coord] higherNeighbours grid here = filter isHigher $ neighbours grid here where this = grid ! here isHigher there = (grid ! there) > this basinSize :: Grid -> Coord -> Int basinSize grid basinSeed = S.size $ breadthFirstSearch grid (S.singleton basinSeed) S.empty breadthFirstSearch :: Grid -> Basin -> Basin -> Basin breadthFirstSearch grid agenda basin | S.null agenda = basin | otherwise = breadthFirstSearch grid agenda' basin' where here = S.findMin agenda candidates = (S.fromList $ higherNeighbours grid here) \\ basin basin' = if (grid ! here) == 9 then basin else S.insert here basin agenda' = S.union candidates $ S.delete here agenda