diff options
Diffstat (limited to 'utils.py')
| -rw-r--r-- | utils.py | 41 |
1 files changed, 40 insertions, 1 deletions
@@ -1,7 +1,8 @@ from __future__ import annotations -from collections import deque +from collections import deque, defaultdict from collections.abc import Iterable, Iterator from dataclasses import dataclass +from heapq import heappop, heappush from itertools import islice from math import sqrt from re import match as re_match @@ -59,6 +60,44 @@ def adjacent_bounded(p: Point2D[int], bound: Point2D[int], diagonal: bool = True -> Iterator[Point2D[int]]: return filter(lambda p: inbounds(p, bound), adjacent(p, diagonal)) +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, intern(start))] + came_from = dict() + g_score = defaultdict(lambda: float('inf')) + g_score[start] = 0 + f_score = {} + f_score[start] = h(start) + while open_set: + while True: + f, current = heappop(open_heapq) + current = itov[current] + if current in open_set: break + if current == goal: + return f + 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) + open_set.add(neighbour) + heappush(open_heapq, (f_score[neighbour], intern(neighbour))) + return None + T = TypeVar('T') def sliding_window(iterable: Iterable[T], n: int) -> Iterator[tuple[T, ...]]: # sliding_window('ABCDEFG', 4) -> ABCD BCDE CDEF DEFG |