diff options
Diffstat (limited to '15.py')
| -rw-r--r-- | 15.py | 54 |
1 files changed, 29 insertions, 25 deletions
@@ -1,21 +1,34 @@ from collections import defaultdict -from utils import open_day +from utils import open_day, Point2D, adjacent_bounded from heapq import heappop, heappush def a_star(start, goal, neighbours, h, d): + vtoi = dict() + itov = dict() + i = 0 + def intern(v): + nonlocal i + ret = vtoi.get(v) + if ret is None: + ret = i + vtoi[v] = i + itov[i] = v + i += 1 + return ret open_set = {start} - open_heapq = [(0, start)] + open_heapq = [(0, intern(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] + while True: + f, current = heappop(open_heapq) + current = itov[current] + if current in open_set: break 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) @@ -23,31 +36,22 @@ def a_star(start, goal, neighbours, h, d): 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)) + open_set.add(neighbour) + heappush(open_heapq, (f_score[neighbour], intern(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) + add = b.x // len(inp[0]) + b.y // len(inp) + x, y = b.x % len(inp[0]), b.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)) +start = Point2D(0, 0) +goal1 = Point2D(len(inp[0]) - 1, len(inp) - 1) +neighbours1 = lambda p: adjacent_bounded(p, Point2D(len(inp[0]), len(inp)), False) +print(a_star(start, goal1, neighbours1, lambda p: 0, d)) +goal2 = Point2D(len(inp[0]) * 5 - 1, len(inp) * 5 - 1) +neighbours2 = lambda p: adjacent_bounded(p, Point2D(len(inp[0]) * 5, len(inp) * 5), False) +print(a_star(start, goal2, neighbours2, lambda p: 0, d)) |