adventofcode2022/lib/astar.py

from heapq import heappush, heappop

def search(start, is_goal, neighbors, heuristic=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 = {}
    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")
            # 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 h is None:
                    continue
                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()