1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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))
|
