1 files changed, 10 insertions, 55 deletions

diff --git a/day12/__init__.py b/day12/__init__.py
index 69f9458..ef03ee0 100644
--- a/day12/__init__.py
+++ b/day12/__init__.py
@@ -5,6 +5,7 @@ from math import inf, sqrt, floor
 from typing import List, Tuple, Iterator, Callable, Any
 
 from aoc import BaseAssignment
+from aoc.mixins import AStarMixin
 
 
 class Map:
@@ -57,7 +58,7 @@ class Map:
         return elevation - 96
 
 
-class Assignment(BaseAssignment, ABC):
+class Assignment(BaseAssignment, AStarMixin[Tuple[int, int]], ABC):
     def distance(self, map: Map, a: Tuple[int, int], b: Tuple[int, int]):
         distance = map.elevation(*b) - map.elevation(*a)
 
@@ -66,67 +67,22 @@ class Assignment(BaseAssignment, ABC):
 
         return 1
 
+
+class AssignmentOne(Assignment):
+    example_result = 31
+
     def heuristic(self, map: Map, a: Tuple[int, int]):
         dx = abs(map.start[0] - a[0])
         dy = abs(map.start[1] - a[1])
         return floor(sqrt(dx**2 + dy**2))
 
-    def a_star(
-        self,
-        map: Map,
-        start: Tuple[int, int],
-        end: Callable[[int, int], bool],
-        distance: Callable[[Tuple[int, int], Tuple[int, int]], int],
-        heuristic: Callable[[Tuple[int, int]], int],
-    ) -> List[Tuple[int, int]]:
-        open = {start}
-        came_from = {}
-
-        g_scores = {start: 0}
-
-        f_scores = {start: heuristic(start)}
-
-        while True:
-            try:
-                current = sorted(
-                    [item for item in open], key=lambda item: f_scores.get(item, inf)
-                )[0]
-            except IndexError:
-                raise RuntimeError("No path found")
-
-            if end(*current):
-                path = [current]
-
-                while current in came_from:
-                    current = came_from[current]
-                    path = [current, *path]
-
-                return path
-
-            open.remove(current)
-
-            for neighbour in map.neighbours(*current):
-                g_score = g_scores.get(current, inf) + distance(current, neighbour)
-
-                if g_score < g_scores.get(neighbour, inf):
-                    came_from[neighbour] = current
-                    g_scores[neighbour] = g_score
-                    f_scores[neighbour] = g_score + heuristic(neighbour)
-
-                    if neighbour not in open:
-                        open.add(neighbour)
-
-
-class AssignmentOne(Assignment):
-    example_result = 31
-
     def run(self, input: Iterator) -> Any:
         map = Map(map=list(input))
 
         path = self.a_star(
-            map,
             map.start,
-            lambda x, y: (x, y) == map.end,
+            lambda coordinate, _: coordinate == map.end,
+            neighbours=lambda a: map.neighbours(*a),
             distance=lambda a, b: self.distance(map, a, b),
             heuristic=lambda a: self.heuristic(map, a),
         )
@@ -141,11 +97,10 @@ class AssignmentTwo(Assignment):
         map = Map(list(input))
 
         path = self.a_star(
-            map,
             map.end,
-            lambda x, y: map.elevation(x, y) == 1,
+            lambda coordinate, _: map.elevation(*coordinate) == 1,
+            neighbours=lambda a: map.neighbours(*a),
             distance=lambda a, b: self.distance(map, b, a),
-            heuristic=lambda a: 1,
         )
 
         return len(path) - 1