a_star.py

def a_star(initial_state: Matrix, end_state: Matrix, heuristic: Callable[[Matrix, Matrix], int] = uniform_cost) -> Node:
    def convert_to_nodes(matrices: List[Matrix]):
        neighbors = []
        for m in matrices:
            key = str(m)
            if key not in graph:
                graph[key] = Node(m, cost=heuristic(m, end_state))
            neighbors.append(graph[key])
        return neighbors

    def update_neighbors(neighbors: List[Node], current: Node):
        for n in neighbors:
            new_depth = current.depth + 1
            if new_depth < n.depth:
                n.parent = current
                n.depth = new_depth
                if n not in to_explore:
                    to_explore.add(n)

    to_explore = PriorityQueue()
    start_node = Node(initial_state, cost=heuristic(initial_state, end_state), depth=0)
    graph: Dict[str, Node] = {str(initial_state): start_node}
    to_explore.add(start_node)
    counter = 0
    while True:
        if to_explore.isempty():
            raise Exception("Failure")
        current = to_explore.pop()
        if current.matrix == end_state:
            logging.info(f'Found min path depth {get_depth(current)} in {counter} tries')
            return current
        adjacent_matrices = get_adjacent_matrices(current.matrix)
        neighbors = convert_to_nodes(adjacent_matrices)
        update_neighbors(neighbors, current)
        counter += 1