diff --git a/graphs/a_star.py b/graphs/a_star.py index 1d7063ccc55a..fd9169f547c8 100644 --- a/graphs/a_star.py +++ b/graphs/a_star.py @@ -1,138 +1,154 @@ +""" +A* (A-Star) Pathfinding Algorithm Implementation. + +A* is an informed search algorithm (heuristic search) that finds the shortest +path between a start node and a goal node in a weighted graph. It balances +the actual distance from the start (g-score) and the estimated distance to +the goal (h-score) using the formula: f(n) = g(n) + h(n). + +Time Complexity: O(E log V) where E is the number of edges and V is the + number of vertices. +Space Complexity: O(V) to store the graph structures and priority queue. +""" + from __future__ import annotations -DIRECTIONS = [ - [-1, 0], # left - [0, -1], # down - [1, 0], # right - [0, 1], # up -] - - -# function to search the path -def search( - grid: list[list[int]], - init: list[int], - goal: list[int], - cost: int, - heuristic: list[list[int]], -) -> tuple[list[list[int]], list[list[int]]]: +import heapq +import math +from collections.abc import Callable + +# ========================================== +# 1. Heuristic Functions +# ========================================== + + +def manhattan_distance(a: tuple[float, float], b: tuple[float, float]) -> float: + """Calculate the Manhattan distance between two 2D points.""" + return abs(a[0] - b[0]) + abs(a[1] - b[1]) + + +def euclidean_distance(a: tuple[float, float], b: tuple[float, float]) -> float: + """Calculate the Euclidean distance between two 2D points.""" + return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2) + + +def chebyshev_distance(a: tuple[float, float], b: tuple[float, float]) -> float: + """Calculate the Chebyshev distance between two 2D points.""" + return max(abs(a[0] - b[0]), abs(a[1] - b[1])) + + +# ========================================== +# 2. Grid-Based Implementation +# ========================================== + + +def a_star_grid( + grid: list[list[float]], + start: tuple[int, int], + end: tuple[int, int], + heuristic_func: Callable[ + [tuple[float, float], tuple[float, float]], float + ] = manhattan_distance, +) -> list[tuple[int, int]] | None: """ - Search for a path on a grid avoiding obstacles. - >>> grid = [[0, 1, 0, 0, 0, 0], - ... [0, 1, 0, 0, 0, 0], - ... [0, 1, 0, 0, 0, 0], - ... [0, 1, 0, 0, 1, 0], - ... [0, 0, 0, 0, 1, 0]] - >>> init = [0, 0] - >>> goal = [len(grid) - 1, len(grid[0]) - 1] - >>> cost = 1 - >>> heuristic = [[0] * len(grid[0]) for _ in range(len(grid))] - >>> heuristic = [[0 for row in range(len(grid[0]))] for col in range(len(grid))] - >>> for i in range(len(grid)): - ... for j in range(len(grid[0])): - ... heuristic[i][j] = abs(i - goal[0]) + abs(j - goal[1]) - ... if grid[i][j] == 1: - ... heuristic[i][j] = 99 - >>> path, action = search(grid, init, goal, cost, heuristic) - >>> path # doctest: +NORMALIZE_WHITESPACE - [[0, 0], [1, 0], [2, 0], [3, 0], [4, 0], [4, 1], [4, 2], [4, 3], [3, 3], - [2, 3], [2, 4], [2, 5], [3, 5], [4, 5]] - >>> action # doctest: +NORMALIZE_WHITESPACE - [[0, 0, 0, 0, 0, 0], [2, 0, 0, 0, 0, 0], [2, 0, 0, 0, 3, 3], - [2, 0, 0, 0, 0, 2], [2, 3, 3, 3, 0, 2]] + Perform A* search on a 2D weighted grid using heapq. + Grid values represent the traversal cost. float('inf') represents an obstacle. + + >>> grid = [[1.0, 1.0, 1.0], [1.0, float('inf'), 1.0], [1.0, 1.0, 1.0]] + >>> a_star_grid(grid, (0, 0), (2, 2), manhattan_distance) + [(0, 0), (0, 1), (0, 2), (1, 2), (2, 2)] """ - closed = [ - [0 for col in range(len(grid[0]))] for row in range(len(grid)) - ] # the reference grid - closed[init[0]][init[1]] = 1 - action = [ - [0 for col in range(len(grid[0]))] for row in range(len(grid)) - ] # the action grid - - x = init[0] - y = init[1] - g = 0 - f = g + heuristic[x][y] # cost from starting cell to destination cell - cell = [[f, g, x, y]] - - found = False # flag that is set when search is complete - resign = False # flag set if we can't find expand - - while not found and not resign: - if len(cell) == 0: - raise ValueError("Algorithm is unable to find solution") - else: # to choose the least costliest action so as to move closer to the goal - cell.sort() - cell.reverse() - next_cell = cell.pop() - x = next_cell[2] - y = next_cell[3] - g = next_cell[1] - - if x == goal[0] and y == goal[1]: - found = True - else: - for i in range(len(DIRECTIONS)): # to try out different valid actions - x2 = x + DIRECTIONS[i][0] - y2 = y + DIRECTIONS[i][1] - if ( - x2 >= 0 - and x2 < len(grid) - and y2 >= 0 - and y2 < len(grid[0]) - and closed[x2][y2] == 0 - and grid[x2][y2] == 0 - ): - g2 = g + cost - f2 = g2 + heuristic[x2][y2] - cell.append([f2, g2, x2, y2]) - closed[x2][y2] = 1 - action[x2][y2] = i - invpath = [] - x = goal[0] - y = goal[1] - invpath.append([x, y]) # we get the reverse path from here - while x != init[0] or y != init[1]: - x2 = x - DIRECTIONS[action[x][y]][0] - y2 = y - DIRECTIONS[action[x][y]][1] - x = x2 - y = y2 - invpath.append([x, y]) - - path = [] - for i in range(len(invpath)): - path.append(invpath[len(invpath) - 1 - i]) - return path, action + rows, cols = len(grid), len(grid[0]) + open_set: list[tuple[float, tuple[int, int]]] = [] + heapq.heappush(open_set, (0.0, start)) + + came_from: dict[tuple[int, int], tuple[int, int]] = {} + g_score = {start: 0.0} + + while open_set: + _, current = heapq.heappop(open_set) + + if current == end: + path = [] + while current in came_from: + path.append(current) + current = came_from[current] + path.append(start) + return path[::-1] + + # 4-directional movement + for move_r, move_c in [(-1, 0), (1, 0), (0, -1), (0, 1)]: + neighbor = (current[0] + move_r, current[1] + move_c) + + if 0 <= neighbor[0] < rows and 0 <= neighbor[1] < cols: + cost = grid[neighbor[0]][neighbor[1]] + if cost == float("inf"): + continue + + tentative_g = g_score[current] + cost + if tentative_g < g_score.get(neighbor, float("inf")): + came_from[neighbor] = current + g_score[neighbor] = tentative_g + f_score = tentative_g + heuristic_func(neighbor, end) + if neighbor not in [item[1] for item in open_set]: + heapq.heappush(open_set, (f_score, neighbor)) + return None + + +# ========================================== +# 3. Adjacency List Implementation +# ========================================== + + +def a_star_adjacency_list( + graph: dict[str, list[tuple[str, float]]], + start: str, + end: str, + heuristic_dict: dict[str, float], +) -> list[str] | None: + """ + Perform A* search on a graph represented as an adjacency list. + heuristic_dict provides pre-calculated h-scores from each node to the goal. + + >>> graph = { + ... 'A': [('B', 1.0), ('C', 4.0)], + ... 'B': [('A', 1.0), ('D', 5.0)], + ... 'C': [('A', 4.0), ('D', 1.0)], + ... 'D': [('B', 5.0), ('C', 1.0)] + ... } + >>> h_dict = {'A': 3.0, 'B': 2.0, 'C': 1.0, 'D': 0.0} + >>> a_star_adjacency_list(graph, 'A', 'D', h_dict) + ['A', 'C', 'D'] + """ + open_set: list[tuple[float, str]] = [] + heapq.heappush(open_set, (0.0, start)) + + came_from: dict[str, str] = {} + g_score = {start: 0.0} + + while open_set: + _, current = heapq.heappop(open_set) + + if current == end: + path = [] + while current in came_from: + path.append(current) + current = came_from[current] + path.append(start) + return path[::-1] + + for neighbor, weight in graph.get(current, []): + tentative_g = g_score[current] + weight + if tentative_g < g_score.get(neighbor, float("inf")): + came_from[neighbor] = current + g_score[neighbor] = tentative_g + f_score = tentative_g + heuristic_dict.get(neighbor, 0.0) + if neighbor not in [item[1] for item in open_set]: + heapq.heappush(open_set, (f_score, neighbor)) + return None if __name__ == "__main__": - grid = [ - [0, 1, 0, 0, 0, 0], - [0, 1, 0, 0, 0, 0], # 0 are free path whereas 1's are obstacles - [0, 1, 0, 0, 0, 0], - [0, 1, 0, 0, 1, 0], - [0, 0, 0, 0, 1, 0], - ] - - init = [0, 0] - # all coordinates are given in format [y,x] - goal = [len(grid) - 1, len(grid[0]) - 1] - cost = 1 - - # the cost map which pushes the path closer to the goal - heuristic = [[0 for row in range(len(grid[0]))] for col in range(len(grid))] - for i in range(len(grid)): - for j in range(len(grid[0])): - heuristic[i][j] = abs(i - goal[0]) + abs(j - goal[1]) - if grid[i][j] == 1: - # added extra penalty in the heuristic map - heuristic[i][j] = 99 - - path, action = search(grid, init, goal, cost, heuristic) - - print("ACTION MAP") - for i in range(len(action)): - print(action[i]) - - for i in range(len(path)): - print(path[i]) + import doctest + + doctest.testmod()