Initial

luogu_P2901_a_star.py

@@ -0,0 +1,124 @@
+import heapq
+import sys
+
+INT_INF = 0x7fffffff
+
+
+class GraphMatrix:
+    def __init__(self, size: int):
+        self.vertices = size
+        self.matrix = [[[] for _ in range(size)] for _ in range(size)]
+
+    def add_edge(self, start: int, end: int, weight: int = 1):
+        self.matrix[start][end].append(weight)
+
+    def get_edge(self, start: int, end: int, invert: bool = False) -> list[int]:
+        if invert:
+            return self.matrix[end][start]
+        else:
+            return self.matrix[start][end]
+
+    def get_edge_min(self, start: int, end: int, invert: bool = False) -> int:
+        return min(self.get_edge(start, end, invert))
+
+    def get_adjacent_vertices(self, vertex: int, invert: bool = False) -> list[int]:
+        return [i for i in range(self.vertices) if self.get_edge(vertex, i, invert)]
+
+    def __repr__(self):
+        return "\n".join([" | ".join(map(str, row)) for row in self.matrix])
+
+
+class VertexNode:
+    def __init__(self, vertex: int, cost: int, heuristic_cost: int = 0):
+        self.vertex = vertex
+        self.cost = cost
+        self.heuristic_cost = heuristic_cost
+
+    @property
+    def evaluated_cost(self):
+        return self.heuristic_cost + self.cost
+
+    def __lt__(self, other):
+        return self.evaluated_cost < other.evaluated_cost
+
+    def __repr__(self):
+        return f"#{self.vertex}(g={self.cost},h={self.heuristic_cost})"
+
+
+def dijkstra_init(graph: GraphMatrix, end: int) -> dict[int, int]:
+    open_list: list[VertexNode] = []
+    heuristic_map: dict[int, int] = {}
+    visited: set[int] = set()
+
+    heuristic_map[end] = 0
+    heapq.heappush(open_list, VertexNode(end, 0))
+    while open_list:
+        current_node = heapq.heappop(open_list)
+        if current_node.vertex in visited:
+            continue
+        visited.add(current_node.vertex)
+
+        for next_vertex in graph.get_adjacent_vertices(current_node.vertex, True):
+            if next_vertex in visited:
+                continue
+            new_cost = current_node.cost + graph.get_edge_min(current_node.vertex, next_vertex, True)
+            # 如果新路径更短，或者该节点第一次被访问
+            if new_cost < heuristic_map.get(next_vertex, INT_INF):
+                heuristic_map[next_vertex] = new_cost
+                heapq.heappush(open_list, VertexNode(next_vertex, new_cost))
+
+    return heuristic_map
+
+def a_star_solve(graph: GraphMatrix,
+                 start: int, end: int,
+                 count: int,
+                 heuristic_map: dict[int, int]) -> list[int]:
+    open_list: list[VertexNode] = []
+    # 此处不是「visited_record」！我们要记录所有访问，只是不需要展开过多。
+    # 因此也不需要「更新 g 值，因为我们记录了一切 g 的节点！
+    # 逻辑是：1. 每条最短路径上的每个节点一定是被 expand 过的；2. 可采用的 h 可保证先找更短路
+    expanded_record: dict[int, int] = {}
+    # node_record: dict[int, VertexNode] = {}
+    original_node = VertexNode(start, 0, heuristic_map.get(start, INT_INF))
+    heapq.heappush(open_list, original_node)
+    # node_record[start] = original_node
+    result_list: list[int] = []
+
+    while open_list:
+        # print(open_list, file=sys.stderr)
+        current_node = heapq.heappop(open_list)
+        expanded_record[current_node.vertex] = expanded_record.get(current_node.vertex, 0) + 1
+        # print(current_node, open_list, file=sys.stderr)
+        if current_node.vertex == end:
+            result_list.append(current_node.cost)
+            if len(result_list) >= count:
+                return result_list
+
+        if expanded_record.get(current_node.vertex, 0) > count:
+            continue  # 不扩展，剪枝
+
+        for next_vertex in graph.get_adjacent_vertices(current_node.vertex):
+            for edge_weight in graph.get_edge(current_node.vertex, next_vertex):
+                next_node = VertexNode(
+                    next_vertex,
+                    current_node.cost + edge_weight,
+                    heuristic_map.get(next_vertex, INT_INF)
+                )
+                heapq.heappush(open_list, next_node)
+
+    return result_list
+
+
+if __name__ == '__main__':
+    vertices_count, edge_count, result_count = map(int, input().split())
+    graph = GraphMatrix(vertices_count + 1)
+    for _ in range(edge_count):
+        start, end, weight = map(int, input().split())
+        graph.add_edge(start, end, weight)
+    # print(graph, file=sys.stderr)
+    heuristic_map = dijkstra_init(graph, 1)
+    print(heuristic_map, file=sys.stderr)
+    result_list = a_star_solve(graph, vertices_count, 1, result_count, heuristic_map)
+    # print(result_list, file=sys.stderr)
+    result_list += [-1] * (result_count - len(result_list))
+    print(*result_list, sep="\n")