91 lines
2.1 KiB
Python
91 lines
2.1 KiB
Python
|
from heapq import heappush, heappop
|
||
|
|
|
||
|
|
def search(start, goal, neighbors, heuristic=None):
|
||
|
|
if heuristic == None:
|
||
|
|
def heuristic(x):
|
||
|
|
return 0
|
||
|
|
# TODO: callable 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 = {}
|
||
|
|
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 this == goal:
|
||
|
|
print("astar: visited", len(done), "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):
|
||
|
|
if n not in done:
|
||
|
|
# calculate the "reduced cost"
|
||
|
|
c = cost_so_far + c + heuristic(n)
|
||
|
|
i += 1
|
||
|
|
heappush(q, (c, 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()
|