Advent of Code 2023 Day 9 - Nx

Mix.install([
  {:kino_aoc, "~> 0.1.5"},
  {:nx, "~> 0.6.4"},
  {:exla, "~> 0.6.4"}
])

Intuition

From the puzzle description, my gut tells me that each line corresponds to a specific polynomial, and making a new sequence from the difference at each step is just calculating derivatives of that polynomial.

For each line, because it can reach all zeroes (i.e. the polynomial becomes a constant function), we know that the highest degree must not be greater than the number of terms minus 1.

For example, as for the sample input 10 13 16 21 30 45, there are 6 terms, and the highest degree must not be greater than 5.

We can write such an polynomial as

$$ y = a_0 x^0 + a_1 x^1 + a_2 x^2 + … + a_n x^n $$

where $$n$$ is the number of terms minus 1.

Now we just need to find all the coefficients $$a_k$$ for all $$k \in [0, n]$$. There are totally $$n$$ coefficients, so we need at least $$n$$ pairs of $$x$$ and $$y$$.

Fortunately, we have enough $$x$$-$$y$$ pair to solve this problem. The $$y$$’s are the numbers in a line, and the $$x$$’s are the indices of the $$y$$’s. I choose 1-based indices, but you can choose any base.

For example, as for the sample input 10 13 16 21 30 45, we have a group of functions:

$$ a_0 \times 1^0 + a_1 \times 1^1 + a_2 \times 1^2 + a_3 \times 1^3 + a_4 \times 1^4 + a_5 \times 1^5 = 10 \ a_0 \times 2^0 + a_1 \times 2^1 + a_2 \times 2^2 + a_3 \times 2^3 + a_4 \times 2^4 + a_5 \times 2^5 = 13 \ a_0 \times 3^0 + a_1 \times 3^1 + a_2 \times 3^2 + a_3 \times 3^3 + a_4 \times 3^4 + a_5 \times 3^5 = 16 \ a_0 \times 4^0 + a_1 \times 4^1 + a_2 \times 4^2 + a_3 \times 4^3 + a_4 \times 4^4 + a_5 \times 4^5 = 21 \ a_0 \times 5^0 + a_1 \times 5^1 + a_2 \times 5^2 + a_3 \times 5^3 + a_4 \times 5^4 + a_5 \times 5^5 = 30 \ a_0 \times 6^0 + a_1 \times 6^1 + a_2 \times 6^2 + a_3 \times 6^3 + a_4 \times 6^4 + a_5 \times 6^5 = 45 $$

If you are familiar with linear algebra, you can quickly convert this group of functions into the matrix form:

$$ \begin{bmatrix} 1^0 & 1^1 & 1^2 & 1^3 & 1^4 & 1^5 \ 2^0 & 2^1 & 2^2 & 2^3 & 2^4 & 2^5 \ 3^0 & 3^1 & 3^2 & 3^3 & 3^4 & 3^5 \ 4^0 & 4^1 & 4^2 & 4^3 & 4^4 & 4^5 \ 5^0 & 5^1 & 5^2 & 5^3 & 5^4 & 5^5 \ 6^0 & 6^1 & 6^2 & 6^3 & 6^4 & 6^5 \ \end{bmatrix} \begin{bmatrix} a_0 \ a_1 \ a_2 \ a_3 \ a_4 \ a_5 \end{bmatrix} = \begin{bmatrix} 10 \ 13 \ 16 \ 21 \ 30 \ 45 \end{bmatrix} $$

Solve this equation, and you got all the coefficients.

After you got all the efficients, then you just pick the next $$x$$, calculate all the powers of it, dot product the coefficients, and you got the answer.

For example, as for the previous example, when you got all the coefficients, you can work out the next number simply by calculating

$$ \begin{bmatrix} 7^0 & 7^1 & 7^2 & 7^3 & 7^4 & 7^5 \end{bmatrix} \begin{bmatrix} a_0 \ a_1 \ a_2 \ a_3 \ a_4 \ a_5 \end{bmatrix} $$

Prep

{:ok, actual_input} =
  KinoAOC.download_puzzle("2023", "9", System.fetch_env!("LB_AOC_SESSION"))

sample_input =
  """
  0 3 6 9 12 15
  1 3 6 10 15 21
  10 13 16 21 30 45
  """
  |> String.trim()

# puzzle_input = sample_input
puzzle_input = actual_input

lines =
  puzzle_input
  |> String.split("\n")
  |> Stream.map(&String.split/1)
  |> Enum.map(&Enum.map(&1, fn s -> String.to_integer(s) end))

Nx.default_backend(EXLA.Backend)

mat_y =
  lines
  |> Nx.tensor(type: :f64)
  |> Nx.transpose()

{terms, _} = Nx.shape(mat_y)

mat_x =
  for x <- 1..terms do
    1 |> Stream.iterate(&amp;(&amp;1 * x)) |> Enum.take(terms)
  end
  |> Nx.tensor(type: :f64)

coeffs = Nx.LinAlg.solve(mat_x, mat_y)

Part 1

vec_x =
  1
  |> Stream.iterate(&(&1 * (terms + 1)))
  |> Enum.take(terms)
  |> Nx.tensor(type: :f64)
  |> Nx.new_axis(0)

vec_x |> Nx.dot(coeffs) |> Nx.sum() |> Nx.to_number()

Part 2

vec_x =
  1
  |> Stream.iterate(&(&1 * 0))
  |> Enum.take(terms)
  |> Nx.tensor(type: :f64)
  |> Nx.new_axis(0)

vec_x |> Nx.dot(coeffs) |> Nx.sum() |> Nx.to_number()