Advent of code day 17

Mix.install([
  {:kino, "~> 0.5.0"}
])

Setup input

example = Kino.Input.textarea("Please paste your input example:")

input = Kino.Input.textarea("Please paste your real input:")

defmodule PriorityQueue do
  @moduledoc """
  A stateless priority queue implementation basing on pairing heap.
  """

  @initial_serial -Bitwise.bsl(1, 64)

  defstruct size: 0,
            serial: @initial_serial,
            heap: nil,
            key_fun: nil

  @typep serial :: integer()
  @typep pairing_heap(key, value) :: nil | {{key, serial()}, value, [pairing_heap(key, value)]}
  @opaque t(key, value) :: %__MODULE__{
            size: non_neg_integer,
            serial: serial(),
            heap: pairing_heap(key, value),
            key_fun: (value -> key)
          }

  @doc """
  Create an empty priority queue.
  """
  @spec new((value -> key)) :: t(key, value) when key: any(), value: any()
  def new(key_fun \\ nil), do: %__MODULE__{key_fun: key_fun || (& &1)}

  @doc """
  Get the size of a priority queue.
  """
  @spec size(t(any, any)) :: non_neg_integer()
  def size(%__MODULE__{size: size}), do: size

  @doc """
  Put an item into a priority queue.
  Returns a new priority queue.
  """
  @spec push(t(key, value), value) :: t(key, value) when key: any(), value: any()
  def push(%__MODULE__{size: size, heap: heap, key_fun: key_fun, serial: serial} = pq, value) do
    %{pq | size: size + 1, heap: meld(heap, to_heap(value, key_fun, serial)), serial: serial + 1}
  end

  @doc """
  Get and delete the smallest item in the priority queue.
  If the queue is empty, returns `{:error, :empty}`,
  otherwise returns `{:ok, smallest_item, new_priority_queue}`.
  """
  @spec pop(t(key, value)) :: {:ok, value, t(key, value)} | {:error, :empty}
        when key: any(), value: any()
  def pop(%__MODULE__{size: 0}), do: {:error, :empty}

  def pop(%__MODULE__{size: size, heap: {_key, value, sub_heaps}} = pq) do
    {:ok, value, %{pq | size: size - 1, heap: pair(sub_heaps)}}
  end

  @doc """
  Get the smallest element in the priority queue.
  """
  @spec top(t(key, value)) :: {:ok, value} | {:error, :empty}
        when key: any(), value: any()
  def top(%__MODULE__{size: 0}), do: {:error, :empty}
  def top(%__MODULE__{heap: {_, value, _}}), do: {:ok, value}

  @doc """
  Merge two priority queues. Crashes if the given priority queues have different key functions.
  """
  @spec merge!(t(key, value), t(key, value)) :: t(key, value)
        when key: any(), value: any()
  def merge!(%__MODULE__{key_fun: f} = pq1, %__MODULE__{key_fun: f} = pq2) do
    %__MODULE__{
      key_fun: f,
      serial: max(pq1.serial, pq2.serial),
      size: pq1.size + pq2.size,
      heap: meld(pq1.heap, pq2.heap)
    }
  end

  def merge!(%__MODULE__{}, %__MODULE__{}) do
    raise ArgumentError, "Can't merge priority queues with different key functions."
  end

  defp to_heap(value, key_fun, serial) do
    {{key_fun.(value), serial}, value, []}
  end

  @spec meld(pairing_heap(key, value), pairing_heap(key, value)) :: pairing_heap(key, value)
        when key: any(), value: any()
  defp meld(nil, heap), do: heap
  defp meld(heap, nil), do: heap
  defp meld({k1, v1, sub1}, {k2, _v2, _sub2} = heap2) when k1 < k2, do: {k1, v1, [heap2 | sub1]}
  defp meld(heap1, {k2, v2, sub2}), do: {k2, v2, [heap1 | sub2]} 

  @spec pair([pairing_heap(key, value)]) :: pairing_heap(key, value)
        when key: any(), value: any()
  defp pair([]), do: nil
  defp pair([node]), do: node 
  defp pair([node1, node2 | rest]), do: meld(meld(node2, node1), pair(rest))
end

defmodule Dijkstra do
  def solve(grid, destination, part) when is_map(grid) do
    pq = PriorityQueue.new() |> PriorityQueue.push({0, {0, 0}, {0, 0}, 0})
    shortest_path(grid, pq, destination, MapSet.new(), part)
  end

  defp shortest_path(grid, pq, destination, seen, part) do
    case PriorityQueue.pop(pq) do
      {:error, :empty} ->
        {:error, :no_path}

      {:ok, {hl, {r, c}, {dr, dc}, n}, pq} ->
        cond do
          {r, c} == destination and can_stop?(n, part) ->
            {:ok, hl}

          MapSet.member?(seen, {r, c, dr, dc, n}) ->
            shortest_path(grid, pq, destination, seen, part)

          true ->
            seen = MapSet.put(seen, {r, c, dr, dc, n})

            pq =
              generate_neighbors(r, c, dr, dc, n, grid, part)
              |> Enum.reduce(pq, fn {new_hl, nr, nc, ndr, ndc, new_n}, acc ->
                PriorityQueue.push(acc, {hl + new_hl, {nr, nc}, {ndr, ndc}, new_n})
              end)

            shortest_path(grid, pq, destination, seen, part)
        end
    end
  end

  defp can_stop?(n, :part1), do: n > 0
  defp can_stop?(n, :part2), do: n >= 4

  defp generate_neighbors(r, c, dr, dc, n, grid, :part2) do
    neighbors =
      if n == 0 do
        for {ndr, ndc} <- [{0, 1}, {1, 0}, {0, -1}, {-1, 0}] do
          {r + ndr, c + ndc, ndr, ndc, 1}
        end
      else
        turns =
          if n >= 4 do
            [
              # Turn left
              {r + dc, c - dr, dc, -dr, 1},
              # Turn right
              {r - dc, c + dr, -dc, dr, 1}
            ]
          else
            []
          end
        if n < 10 do
          [{r + dr, c + dc, dr, dc, n + 1} | turns]
        else
          turns
        end
      end

    for {nr, nc, ndr, ndc, new_n} <- neighbors,
        (new_cost = grid[{nr, nc}]) != nil do
      {new_cost, nr, nc, ndr, ndc, new_n}
    end
  end

  defp generate_neighbors(r, c, dr, dc, n, grid, :part1) do
    # Continue straight if n < 3 and not at starting position
    straight =
      if n < 3 and {dr, dc} != {0, 0} do
        nr = r + dr
        nc = c + dc

        case grid[{nr, nc}] do
          nil -> []
          cost -> [{cost, nr, nc, dr, dc, n + 1}]
        end
      else
        []
      end

    turns =
      for {ndr, ndc} <- [{0, 1}, {1, 0}, {0, -1}, {-1, 0}],
          {ndr, ndc} != {dr, dc},
          {ndr, ndc} != {-dr, -dc},
          nr = r + ndr,
          nc = c + ndc,
          cost = grid[{nr, nc}],
          cost != nil do
        {cost, nr, nc, ndr, ndc, 1}
      end

    straight ++ turns
  end
end

parsed =
  example
  |> Kino.Input.read()
  |> String.split("\n", trim: true)
  |> Enum.map(&(String.split(&1, "", trim: true) |> List.to_tuple()))
  |> List.to_tuple()

lines = tuple_size(parsed) - 1
cols = tuple_size(elem(parsed, 0)) - 1

grid =
  for l <- 0..lines, c <- 0..cols, into: %{} do
    {{l, c}, elem(elem(parsed, l), c) |> String.to_integer()}
  end

Part 01

Dijkstra.solve(grid, {lines, cols}, :part1)

Part 02

Dijkstra.solve(grid, {lines, cols}, :part2)