leah blogs

December 2021

07dec2021 · Counting lanternfish with BQN and linear algebra

Yesterday, Advent of Code had an interesting problem: modeling the growth rate of lanternfish. Read the link for the full details, but the core algorithm is this:

Each day, a 0 becomes a 6 and adds a new 8 to the end of the list, while each other number decreases by 1 if it was present at the start of the day.

For part 1, you need to compute the total number of lanternfish after 80 days, and for part 2 after 256 days. While the first part is possible with a native list based version (around 300k), the second part yields over 1.6 trillion for my input, so storing every lanternfish on its own is infeasible.

Luckily, with a bit of thinking we quickly realize that the position and order of the lanternfish is irrelevant and we can thus just count how many at each age we have and sum them up at the end for the answer.

I decided to tackle this problem using the array language BQN and will try to explain how it works.

First, let’s convert the sample input data into a histogram: For reasons that will be obvious in a minute, we compare each input element with the list of numbers from 0 to 8 (↕9).

   input ← ⟨3,4,3,1,2⟩
   input =⌜ ↕9
┌─
╵ 0 0 0 1 0 0 0 0 0
  0 0 0 0 1 0 0 0 0
  0 0 0 1 0 0 0 0 0
  0 1 0 0 0 0 0 0 0
  0 0 1 0 0 0 0 0 0
                    ┘

If we sum this table by column now, we have a count which age appears how often (note that ages over 8 never happen):

   d ← +˝ input =⌜ ↕9
⟨ 0 1 1 2 1 0 0 0 0 ⟩

Next, we need to figure out what happens in a step. We can model one part (number decreases by 1, or becomes 8 when it was 0) by rotating () the array to the left:

   1 ⌽ d
⟨ 1 1 2 1 0 0 0 0 0 ⟩
   2 ⌽ d
⟨ 1 2 1 0 0 0 0 0 1 ⟩

But we also need to add the first element to column 6, e.g. by adding an appropriately scaled vector. First, the vector, which we get by taking the numbers of 0 to 8 and comparing them to 6:

   6=↕9
⟨ 0 0 0 0 0 0 1 0 0 ⟩

Now we can scale this vector by the first element ():

   (6=↕9) × ⊑ ⟨9,8,7,6,5,4,3,2,1⟩
⟨ 0 0 0 0 0 0 9 0 0 ⟩

Finally, we put both things together. We can use a hook if we bind () the argument of the rotation to the operator:

Step ← 1⊸⌽ + (6=↕9)×⊑

Now, we simply need to repeat () this step 80 (resp. 256) times and count up the results to solve the puzzle!

   +´ Step⍟80 d
5934
   +´ Step⍟256 d
26984457539

This agrees with the provided example.

Of course, computing 256 steps is very fast, but we can wonder if there is not a more efficient solution. Let’s say we wanted to compute many lanternfish populations quickly.

The key realization is that Step is a linear function, which means:

(n × Step v) ≡ (Step n × v)
(Step v) + (Step w) ≡ (Step v + w)

We can thus compute Step for every basis vector of our 9-dimensional histogram space; let’s take the unit vectors here:

   step256 ← {+´Step⍟256 𝕩}˘ =⌜˜↕9
⟨ 6703087164 6206821033 5617089148 5217223242 4726100874
  4368232009 3989468462 3649885552 3369186778 ⟩

And now we can solve part 2 for any input by using a dot-product (pairwise multiplication followed by addition):

   step256 +´∘× d
26984457539

This yields the same result as above, but only does 9 multiplications and 9 additions.

We can also compute step256 faster by using a bit of matrix theory. Since we know the operation of step on the basis vectors, we can compute the matrix m corresponding to this linear operator:

   m ← Step˘ =⌜˜↕9
┌─
╵ 0 0 0 0 0 0 1 0 1
  1 0 0 0 0 0 0 0 0
  0 1 0 0 0 0 0 0 0
  0 0 1 0 0 0 0 0 0
  0 0 0 1 0 0 0 0 0
  0 0 0 0 1 0 0 0 0
  0 0 0 0 0 1 0 0 0
  0 0 0 0 0 0 1 0 0
  0 0 0 0 0 0 0 1 0
                    ┘

Note how we just applied Step for every column of the identity matrix.

We can solve the problem the same way now by computing the 256th power of the matrix and multiplying it to the histogram vector and summing up again.

First, we need the matrix product:

   MP ← +˝∘×⎉1‿∞

Naively, we can compute the power like this (we already have m to the power of 1!):

   m256 ← m MP⍟255 m

Now multiply d to it, and sum up:

   +´ m256 +˝∘× d
26984457539

Note that step256 is just the sum of rows, so we can precompute that:

   step256 ≡ +˝˘ m256
1

If you paid attention, you may now wonder why we replaced 256 row operations with 255 matrix multiplications and claim it’s faster to compute. We need an additional trick: fast matrix exponentiation. Instead of 255 matrix multiplications we just need 8 matrix squarings with itself to compute the 256th power:

   m256 ≡ MP˜⍟8 m
1

This yields an asymptotically faster algorithm to compute the step vector for higher powers. For non-powers of two you need to implement a square-and-multiply algorithm, or look up the optimal result. This is left as an exercise to the reader.

For day 80 we can use only 7 matrix multiplications:

   (MP˜⍟4 (m MP MP˜⍟2 m)) ≡ (m MP⍟79 m)
1

In theory, you could optimize this even further and compute a closed form from the eigenvectors of m, but you’ll end up with nonic roots and I have not found a way to compute them precise enough to be useful. So let’s leave it at that for now.

NP: Slade—My Oh My

Copyright © 2004–2021