Part 1

Today we’re given a 2-D matrix of natural numbers (a heightmap) and must first find the “low points”. A low point is defined to be an entry which is less than each of its four row-and-column neighbors.

Our “map” will thus support a 2-D notion of indexing, as well as a map from each index to its neighbors:

∀ m, n : ℕ → Height_Map (m, n) ≅ [(i, j) | i ← [0..m-1], j ← [0..n-1]]

elem_at : Height_Map (m, n) → (ℕ × ℕ) → ℕ

neighbors : Height_Map (m, n) → (ℕ × ℕ) → [ℕ]

For simplicity, we assume all indices are within bounds.

neighbors gives the neighbors of the element at i, j in Height_Map; its value will be a list of two, three, or four numbers, depending on where i, j is in the map (corner elements have only two neighbors, etc.). We assume indices are row-major:

neighbors hm = list (elem_at hm) ∘ neighbor_ixs (bounds hm)

neighbor_ixs : (ℕ × ℕ) → (ℕ × ℕ) → [(ℕ × ℕ)]
neighbor_ixs (m, n) (r, c) =
  filter (≠ (r, c))
    (nub [(r ∸ 1, c), (r, c ∸ 1), (r, min (c + 1) n),
         (min (r + 1) m, c)])

neighbor_ixs is a little tedious. ∸ here is natural number subtraction (AKA “monus”), which returns zero if the first argument is less than the second. nub : Eq α → [α] → [α] is a library function which removes duplicate equal elements from a list.

We can now specify a function that produces the low points of a Height_Map.

low : Height_Map (m, n) → (ℕ × ℕ) → Bool
low hm (i, j) =
  let (e, ns) = (elem_at hm (i, j), neighbors hm (i, j)) in
    all (λ n → e < n) ns

low_points : Height_Map (m, n) → [ℕ × ℕ]
low_points hm = [ (i, j) | i ← [0..m-1],
                           j ← [0..n-1],
                           low hm (i, j) ]

low_point_heights hm = list (elem_at hm) ∘ low_points hm

The rest is easy. We want the sum of the “risk levels” of the low points:

risk : ℕ → ℕ
risk n = n + 1

The whole thing is thus

sum ∘ list risk ∘ low_points_heights

As usual, this can be fused into a single traversal, but it doesn’t seem worth the trouble.

Executable Haskell solution

Haskellized short test input

Part 2

Now things get much more interesting. We want to find basins; a basin is described informally as a collection of “all locations that eventually flow downward to a single low point”. Locations of height 9 are not included in any basin.

Let’s try to come up with a more formal definition.

This immediately suggests an algorithm for finding the basin associated with a low point p. We can see each location as a node of a graph with four edges linking it to its neighbors; to find a basin, we begin with just a low point, add the neighboring nodes that satisfy the above conditions, then recurse on each neighbor.

This allows us to construct a tree (specifically, a rose tree) describing a basin. Each node associates a location in the basin with its “watershed” (the node’s children), so we can think of this as a model of the basin’s “flow”:

data Rose α   = Node α (Forest α)
data Forest α = [Rose α]

basin_flow : Height_Map (m, n) → (ℕ × ℕ) → Rose (ℕ × ℕ)
basin_flow hm (i, j) =
  let h          = elem_at hm (i, j)
      ns         = neighbor_ixs (m, n) (i, j)
    in Node (i, j)
            (Forest (list (basin_flow hm)
                          (filter (inb h ∘ elem_at hm) ns)))

inb : ℕ → ℕ → Bool
inb h k = k ≠ 9 ∧ k > h

Rose trees have the following mutually-recursive pairs of fold and unfold functions (from Gibbons 2003):

foldR : (α → γ → β) → ([β] → γ) → Rose α → β
foldR f g (Node x ts) = f x (foldF f g ts)

foldF : (α → γ → β) → ([β] → γ) → Forest α → γ
foldF f g = g ∘ list (foldR f g)

unfoldR : (β → α) → (β → [β]) → β → Rose α
unfoldR f g x = Node (f x) (unfoldF f g x)

unfoldF : (β → α) → (β → [β]) → β → Forest α
unfoldF f g = list (unfoldR f g) ∘ g

Reworking basin_flow slightly, we can express it as a rose tree unfold:

basin_flow : Height_Map (m, n) → (ℕ × ℕ) → Rose (ℕ × ℕ)
basin_flow hm = unfoldR id (expand hm)

expand : Height_Map (m, n) → ((ℕ × ℕ) → [(ℕ × ℕ)])
expand hm (i, j) =
  let h = elem_at hm (i, j) in
    [ (k, l) | (k, l) ← neighbor_ixs (m, n) (i, j),
               inb h (elem_at hm (k, l)) ]

Note the use of use of curried application here to keep the height map separate from the unfold. All of the details here are in expand, which expands the basin by a step, that is, it generates the surrounding set of basin locations for a given point.

Every low point in our map is part of exactly one basin (since two basins sharing the same low point would be contiguous at that low point, and hence we’d have only one basin). Thus, we can express the list of all basins in our map by:

basins hm = list (nub ∘ flatten ∘ basin_flow hm) (low_points hm)

flatten : Rose α → [α]
flatten = foldR (∷) concat

The library function nub : [α] → [α] is used to remove duplicate nodes, which can be introduced by our method for finding basins; it may happen that two nodes share the same neighbor, which is of greater height than both.

The list of sizes is then

basin_sizes hm = list length (basins hm)

Using the functor properties of list, this becomes

basin_sizes hm = list (length ∘ nub ∘ flatten ∘ basin_flow hm)
                      (low_points hm)

In practice, this definition is quite efficient enough for solving the current puzzle. However, let’s improve it. As many have probably already seen, we don’t need to construct the intermediate “flow trees” at all. Using a quick calculation, we derive a hylomorphism operator for rose trees that allows us to “deforest” our algorithm.

We specify:

hyloR :: (α → γ → δ) → ([δ] → γ) → (β → α) → (β → [β]) → β → δ
hyloR f g h l = foldR f g ∘ unfoldR h l

The calculation is straightforward. For all x : β,

  hyloR f g h l x

=  { specification }

  foldR f g (unfoldR h l x)

=  { unfoldR.1 }

  foldR f g (Node (h x) (unfoldF h l x))

=  { foldR.1 }

  f (h x) (foldF f g (unfoldF h l x))

=  { specify: hyloF f g h l = foldF f g ∘ unfoldF h l }

  f (h x) (hyloF f g h l x)

As expected, hyloR is one of a pair of mutually-recursive functions. Now we calculate its Forest counterpart, hyloF:

  hyloF f g h l x

=  { specification }

  foldF f g (unfoldF h l x)

=  { unfoldF.1 }

  foldF f g (list (unfoldR h l) (l x))

=  { foldF.1 }

  g (list (foldR f g) (list (unfoldR h l) (l x)))

=  { functor property of list }

  g (list (foldR f g ∘ unfoldR h l) (l x))

=  { specification of hyloR }

  g (list (hyloR f g h l) (l x))

This completes the derivation of hyloR and hyloF.

Since for all hm : Height_Map (m, n) we have,

flatten ∘ basin_flow hm = foldR (∷) concat ∘ unfoldR id (expand hm)

we now redefine

basin_sizes hm = list (length ∘ nub ∘ enum_basins hm)
                      (low_points hm)

enum_basins : Height_Map (m, n) → (ℕ × ℕ) → [(ℕ × ℕ)]
enum_basins hm = hyloR (∷) concat id (expand hm)

enum_basins is thus a function which, given a low point, produces the list of the coordinates of the surrounding basin.

A further set of calculations would allow us to fuse length, nub, and enum_basins, eliminating the intermediate lists, as well. The fusion of length and nub doesn’t produce a significant improvement, however, and fusing the resulting list catamorphism with the rose tree hylomorphism enum_basins requires quite a bit more work. We’ve eliminated the trees, so, for now, we’ll call this enough.

The rest of the puzzle is bookkeeping. We want to know the product of the sizes of the three largest basins:

three_biggest = take 3 ∘ sort_decreasing

solution hm = product ∘ three_biggest ∘ basin_sizes hm

Executable Haskell solution


Gibbons, Jeremy & de Moor, Oege, eds. (2003). The Fun of Programming. Palmgrave Macmillan. See Chapter 3, “Origami Programming”.

AOC 2021 Index