Advent of Code 2021 day 9

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 bounds and 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]
  ]

Part 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

Part 2

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

Code

You can get the code from my locally-hosed Git repo, or from Gitlab.