125 lines
4.7 KiB
Python
125 lines
4.7 KiB
Python
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")
|