Day 4

Mix.install([:exla])

Nx.default_backend(EXLA.Backend)

Section

defmodule Day4 do
  import Nx.Defn

  def part_1(input) do
    # Using Nx just for fun. Could be done with MapSet :p
    input
    |> String.split("\n", trim: true)
    |> Enum.map(fn line ->
      [elf1_str, elf2_str] = String.split(line, ",", trim: true)
      [e1_s, e1_e] = elf1_str |> String.split("-") |> Enum.map(&String.to_integer/1)
      [e2_s, e2_e] = elf2_str |> String.split("-") |> Enum.map(&String.to_integer/1)

      template_last_idx = max(e1_e, e2_e)

      {Nx.broadcast(0, {template_last_idx + 1}), Nx.tensor(Enum.to_list(e1_s..e1_e)),
       Nx.tensor(Enum.to_list(e2_s..e2_e))}
    end)
    |> Enum.map(&calculate_full_overlap/1)
    |> Nx.stack()
    |> Nx.sum()
  end

  defn calculate_full_overlap({target, idx1, idx2}) do
    t1 = Nx.indexed_add(target, Nx.new_axis(idx1, 1), Nx.broadcast(1, idx1))
    t2 = Nx.indexed_add(target, Nx.new_axis(idx2, 1), Nx.broadcast(1, idx2))

    overlaps = t1 and t2

    Nx.all(overlaps == t1) or Nx.all(overlaps == t2)
  end

  def part_2(input) do
    # Using Nx just for fun. Could be done with MapSet :p
    input
    |> String.split("\n", trim: true)
    |> Enum.map(fn line ->
      [elf1_str, elf2_str] = String.split(line, ",", trim: true)
      [e1_s, e1_e] = elf1_str |> String.split("-") |> Enum.map(&String.to_integer/1)
      [e2_s, e2_e] = elf2_str |> String.split("-") |> Enum.map(&String.to_integer/1)

      template_last_idx = max(e1_e, e2_e)

      {Nx.broadcast(0, {template_last_idx + 1}), Nx.tensor(Enum.to_list(e1_s..e1_e)),
       Nx.tensor(Enum.to_list(e2_s..e2_e))}
    end)
    |> Enum.map(fn {target, idx1, idx2} ->
      target
      |> Nx.indexed_add(Nx.new_axis(idx1, 1), Nx.broadcast(1, idx1))
      |> Nx.indexed_add(Nx.new_axis(idx2, 1), Nx.broadcast(1, idx2))
      |> Nx.greater(1)
      |> Nx.any()
    end)
    |> Nx.stack()
    |> Nx.sum()
  end
end

test_input = """
2-4,6-8
2-3,4-5
5-7,7-9
2-8,3-7
6-6,4-6
2-6,4-8
"""

Day4.part_1(test_input)

input = File.read!("day4_input.txt")

Day4.part_1(input)

Day4.part_2(test_input)

Day4.part_2(input)