advent_of_code/elixir/livebook/2024/day13.livemd

AOC 2024 Day 13

Mix.install([
  {:nx, "0.9.2"}
])

Section

input = """
"""

defmodule ClawContraption do
  def part1(input), do: parse(input, 0) |> solve()
  def part2(input), do: parse(input, 10000000000000) |> solve()
  
  def parse(input, n) do
    (for line <- (input |> String.trim() |> String.split("\n")), reduce: {[], []} do
      {vecs, matrices} = acc -> 
       case line do
         "Button " <> rest ->
           [[_, a, b]] = Regex.scan(~r/[AB]: X\+(\d+), Y\+(\d+)/, rest)
           v = [a, b] |> Enum.map(&String.to_integer(&1))
           {[v | vecs], matrices}
         "Prize: " <> rest ->
           [[_, a, b]] = Regex.scan(~r/X=(\d+), Y=(\d+)/, rest)
           v = [a, b] |> Enum.map(&String.to_integer(&1))
           a = Nx.tensor(vecs, type: :f64) |> Nx.transpose()
           b = Nx.tensor(v |> Enum.map(&(&1+n)), type: :f64)
           {[], [{a, b} | matrices]}
         _ -> acc
       end
        end |> elem(1))
  end
  
  def solve(systems) do
    for {a, b} <- systems,
        c = Nx.LinAlg.invert(a) |> Nx.dot(b)  |> Nx.round(),
        b_prime = Nx.dot(
            a, c
        ) |> Nx.round() |> Nx.as_type(:f64),
      c = Nx.as_type(c, :s64),
        b == b_prime,
      reduce: Nx.tensor(0, type: :s64)
      do
        acc2 ->
          c |> Nx.dot(Nx.tensor([1, 3], type: :s64)) |> Nx.add(acc2)
      end 
  end
end

ClawContraption.part1(input)

ClawContraption.part2(input)