Day 17

lib/day17.ex

@@ -0,0 +1,167 @@
+defmodule Aoc2023.Day17 do
+  use Aoc2023
+
+  defmodule PriorityQueue do
+    defstruct [:set]
+    def new(contents \\ []), do: %__MODULE__{set: :gb_sets.from_list(contents)}
+
+    def add(%__MODULE__{set: set}, elem, prio) do
+      %__MODULE__{set: :gb_sets.add({prio, elem}, set)}
+    end
+
+    def size(%__MODULE__{set: set}) do
+      :gb_sets.size(set)
+    end
+
+    def pop_min(%__MODULE__{set: set}) do
+      case :gb_sets.size(set) do
+        0 ->
+          {nil, new()}
+
+        _ ->
+          {res, set} = :gb_sets.take_smallest(set)
+          {res, %__MODULE__{set: set}}
+      end
+    end
+
+    defimpl Inspect do
+      import Inspect.Algebra
+
+      def inspect(%PriorityQueue{set: set}, opts) do
+        concat(["#PriorityQueue.new(", to_doc(:gb_sets.to_list(set), opts), ")"])
+      end
+    end
+  end
+
+  def parse(input) do
+    lines = input |> lines
+
+    {
+      {lines |> length() |> Kernel.-(1), lines |> head |> String.length() |> Kernel.-(1)},
+      lines
+      |> Enum.with_index()
+      |> Enum.flat_map(fn {row, y} ->
+        row
+        |> String.to_charlist()
+        |> Enum.with_index()
+        |> Enum.map(fn {cost, x} -> {{y, x}, String.to_integer(<<cost>>)} end)
+      end)
+      |> Enum.into(%{})
+    }
+  end
+
+  defp a_star(starts, neighbors, heuristic, is_goal) do
+    loop = fn recur, explored, frontier ->
+      {explored, frontier}
+      {{_, current}, frontier} = PriorityQueue.pop_min(frontier)
+
+      if is_goal.(current) do
+        {explored, current}
+      else
+        cost_to_here = Map.get(explored, current) |> elem(1)
+
+        {explored, frontier} =
+          for {neighbor, cost} <- neighbors.(current), reduce: {explored, frontier} do
+            {explored, frontier} ->
+              cost_ = cost_to_here + cost
+
+              if Map.get(explored, neighbor, {nil, cost_}) |> elem(1) < cost_ do
+                {explored, frontier}
+              else
+                {
+                  Map.put(explored, neighbor, {current, cost_}),
+                  PriorityQueue.add(frontier, neighbor, cost_ + heuristic.(neighbor))
+                }
+              end
+          end
+
+        recur.(recur, explored, frontier)
+      end
+    end
+
+    {explored, goal} =
+      loop.(
+        loop,
+        starts |> Enum.map(&{&1, {nil, 0}}) |> Enum.into(%{}),
+        starts |> Enum.map(&{0, &1}) |> PriorityQueue.new()
+      )
+
+    # This doesn't include the starting point
+    path_find = fn recur, point, path ->
+      case Map.get(explored, point) do
+        nil -> path
+        {nil, _} -> path
+        {prev, _} -> recur.(recur, prev, [point | path])
+      end
+    end
+
+    path_find.(path_find, goal, [])
+  end
+
+  def part1({{my, mx}, grid}) do
+    # Nodes are y, x, dir, straight length
+    a_star(
+      [{0, 0, :d, 0}, {0, 0, :r, 0}],
+      fn {y, x, dir, len} ->
+        case dir do
+          :u ->
+            [{y, x - 1, :l, 0}, {y, x + 1, :r, 0}] ++
+              if len < 2, do: [{y + 1, x, dir, len + 1}], else: []
+
+          :d ->
+            [{y, x - 1, :l, 0}, {y, x + 1, :r, 0}] ++
+              if len < 2, do: [{y - 1, x, dir, len + 1}], else: []
+
+          :l ->
+            [{y + 1, x, :u, 0}, {y - 1, x, :d, 0}] ++
+              if len < 2, do: [{y, x - 1, dir, len + 1}], else: []
+
+          :r ->
+            [{y + 1, x, :u, 0}, {y - 1, x, :d, 0}] ++
+              if len < 2, do: [{y, x + 1, dir, len + 1}], else: []
+        end
+        |> Enum.filter(fn {y, x, _, _} -> y in 0..my and x in 0..mx end)
+        |> Enum.map(fn {y, x, _, _} = pos -> {pos, Map.get(grid, {y, x})} end)
+      end,
+      fn {y, x, _, _} -> abs(y - my) + abs(x - mx) end,
+      fn {y, x, _, _} -> {y, x} == {my, mx} end
+    )
+    |> Enum.map(fn {y, x, _, _} -> Map.get(grid, {y, x}) end)
+    |> Enum.sum()
+  end
+
+  def part2({{my, mx}, grid}) do
+    # Nodes are y, x, dir, straight length
+    a_star(
+      [{0, 0, :r, 0}, {0, 0, :d, 0}, {0, 0, :u, 0}, {0, 0, :l, 0}],
+      fn {y, x, dir, len} ->
+        (if(len < 9,
+           do:
+             case dir do
+               :u -> [{y + 1, x, :u, len + 1}]
+               :d -> [{y - 1, x, :d, len + 1}]
+               :l -> [{y, x - 1, :l, len + 1}]
+               :r -> [{y, x + 1, :r, len + 1}]
+             end,
+           else: []
+         ) ++
+           if(len > 2,
+             do:
+               case dir do
+                 :u -> [{y, x - 1, :l, 0}, {y, x + 1, :r, 0}]
+                 :d -> [{y, x - 1, :l, 0}, {y, x + 1, :r, 0}]
+                 :l -> [{y + 1, x, :u, 0}, {y - 1, x, :d, 0}]
+                 :r -> [{y + 1, x, :u, 0}, {y - 1, x, :d, 0}]
+               end,
+             else: []
+           ))
+        |> Enum.filter(fn {y, x, _, _} -> y in 0..my and x in 0..mx end)
+        |> Enum.map(fn {y, x, _, _} = pos -> {pos, Map.get(grid, {y, x})} end)
+      end,
+      fn {y, x, _, _} -> abs(y - my) + abs(x - mx) end,
+      fn {y, x, _, len} -> {y, x} == {my, mx} and len > 2 end
+    )
+    |> Enum.map(fn {y, x, _, _} -> Map.get(grid, {y, x}) end)
+    |> Enum.sum()
+  end
+end