diff options
| -rw-r--r-- | 15.py | 53 |
1 files changed, 53 insertions, 0 deletions
@@ -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)) |