adventofcode2022/lib/astar.py

106 lines
2.7 KiB
Python

from heapq import heappush, heappop
def search(start, is_goal, neighbors, heuristic=None, worst=None):
if heuristic == None:
def heuristic(x):
return 0
if not callable(is_goal):
goal = is_goal
def is_goal(this):
return this == goal
if not isinstance(start, list):
start = [start]
i = 0 # tiebreaker
q = []
for s in start:
i += 1
heappush(q, (heuristic(s), i, s, None))
done = {}
if worst:
best_worst = min(worst(s) for s in start)
while q:
z, _, this, prev = heappop(q)
if this in done:
continue
cost_so_far = z - heuristic(this)
#print(this,z, cost_so_far)
done[this] = (cost_so_far, prev)
if is_goal(this):
print("astar: visited", len(done), "nodes")
print("astar: pending", len(q), "nodes")
# reconsruct the path
n = this
path = []
while n is not None:
path.append(n)
_, n = done[n]
path.reverse()
return cost_so_far, this, path
for c, n in neighbors(this):
c = cost_so_far + c
if n not in done or c < done[n][0]:
h = heuristic(n)
if worst:
if c+h > best_worst:
# if the best possible cost for this node
# is worse than the lowest worst-case cost we've seen
# then don't even bother exploring it
continue
w = worst(n)
if c+w < best_worst:
best_worst = c+w
i += 1
heappush(q, (c+h, i, n, this))
return float('inf'), None, []
def test():
data = [
"aabqponm",
"abcryxxl",
"accszzxk",
"acctuvwj",
"abdefghi",
]
#data = [
# "aabbaaba",
# "abcbaaba",
# "accabcbb",
# "accbbbba",
# "abbbabab",
#]
start = (0,0)
goal = (5,2)
def get(x,y):
return ord(data[y][x])
def neighbors(src):
x,y = src
here = get(x,y)
n = []
def push(x,y):
if 0 <= y < len(data):
if 0 <= x < len(data[y]):
if get(x,y) <= here+1:
n.append((1, (x,y)))
push(x-1, y)
push(x+1,y)
push(x, y-1)
push(x, y+1)
return n
def heuristic(n):
return abs(goal[0] - n[0]) + abs(goal[1] - n[1])
c, _, path = search(start, goal, neighbors, heuristic)
print(*data, sep="\n")
print(*path, sep="\n")
assert c == 31, c
if __name__ == '__main__':
test()