summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--15.py53
1 files changed, 53 insertions, 0 deletions
diff --git a/15.py b/15.py
new file mode 100644
index 0000000..ab369f5
--- /dev/null
+++ b/15.py
@@ -0,0 +1,53 @@
+from collections import defaultdict
+from utils import open_day
+from heapq import heappop, heappush
+
+def a_star(start, goal, neighbours, h, d):
+ open_set = {start}
+ open_heapq = [(0, start)]
+ came_from = dict()
+ g_score = defaultdict(lambda: float('inf'))
+ g_score[start] = 0
+ f_score = defaultdict(lambda: float('inf'))
+ f_score[start] = h(start)
+
+ while open_set:
+ f, current = open_heapq[0]
+ if current == goal:
+ return f
+ heappop(open_heapq)
+ open_set.remove(current)
+ for neighbour in neighbours(current):
+ tentative_g = g_score[current] + d(current, neighbour)
+ if tentative_g < g_score[neighbour]:
+ came_from[neighbour] = current
+ g_score[neighbour] = tentative_g
+ f_score[neighbour] = tentative_g + h(neighbour)
+ if neighbour not in open_set:
+ open_set.add(neighbour)
+ heappush(open_heapq, (f_score[neighbour], neighbour))
+ return None
+
+with open_day(15) as f:
+ inp = [[int(c) for c in l.rstrip()] for l in f]
+
+def neighbours(p, dim):
+ for dx, dy in ((-1, 0), (1, 0), (0, -1), (0, 1)):
+ px, py = p[0] + dx, p[1] + dy
+ if px < 0 or py < 0 or px >= dim[0] or py >= dim[1]: continue
+ yield px, py
+
+def manhattan(a, b): return abs(a[0] - b[0]) + abs(a[1] - b[1])
+def d(a, b):
+ x, y = b
+ add = x // len(inp[0]) + y // len(inp)
+ x, y = x % len(inp[0]), y % len(inp)
+ return (inp[y][x] + add - 1) % 9 + 1
+
+start = (0, 0)
+goal1 = (len(inp[0]) - 1, len(inp) - 1)
+neighbours1 = lambda p: neighbours(p, (len(inp[0]), len(inp)))
+print(a_star(start, goal1, neighbours1, lambda p: manhattan(p, goal1), d))
+goal2 = (len(inp[0]) * 5 - 1, len(inp) * 5 - 1)
+neighbours2 = lambda p: neighbours(p, (len(inp[0]) * 5, len(inp) * 5))
+print(a_star(start, goal2, neighbours2, lambda p: manhattan(p, goal2), d))