AcWing A*算法相关问题 178. 第K短路

import sys
sys.stdin = open('data.txt', 'r')


'''
先利用反向边用单元最短路求所有点到终点的最短距离，把这个距离作为预估值做A*，第k次出
队列的状态中的距离就是答案

'''

from typing import List, Tuple
from queue import PriorityQueue

class DijkStra:
    # start_node 是单源最短路起点 edges是边的三元组(节点1， 节点2， 边权重) nodes是所有节点的列表
    def __init__(self, start_node, edges: List[Tuple[int, int, int]], nodes: List[int], is_directed = False):
        self.edges = edges
        self.nodes = nodes
        self.start_node = start_node
        self.is_directed = is_directed

    def __init(self):
        edges = self.edges
        nodes = self.nodes
        start_node = self.start_node

        self.link = {node: {} for node in nodes}  # 邻接表
        self.start_node = start_node
        self.S = {}  # 已经确定最短路径长度的节点集合
        self.total_nodes_num = len(nodes)
        U = {node: 0x7fffffff for node in nodes if node != start_node}  # 还未确定最短路径长度的节点集合
        U[start_node] = 0

        for a, b, edge_len in edges:
            self.link[a][b] = edge_len
            if not self.is_directed:
                self.link[b][a] = edge_len

            if a == start_node:
                U[b] = edge_len

            if not self.is_directed:
                if b == start_node:
                    U[a] = edge_len

        self.U_que = PriorityQueue()  # 维护U集合的小顶堆，堆里面的内容是(路径长度， 节点值)的二元组
        for k, v in U.items():
            self.U_que.put((v, k))


    # 获取最短路径长度的列表[(节点1，长度1), (节点2, 长度2) .....]
    def getMinPathLen(self) -> List[Tuple[int, int]]:
        return self.getMinInfo(False)

    # 获取所有最短路径列表[(节点1, 长度1, 路径节点列表1)， (节点2, 长度2, 路径节点列表2)]
    def getMinPath(self):
        return self.getMinInfo(True)

    def getMinInfo(self, get_path):
        self.__init()
        pre = {}  # 记录每个节点结尾的最短路径的前驱节点

        while len(self.S) != self.total_nodes_num:
            while True:
                min_edge_len, min_node = self.U_que.get_nowait()  # 从U中把当前路径长度最小的节点取出来
                if min_node not in self.S:
                    break

            self.S[min_node] = min_edge_len

            # 更新U集合中节点的距离
            for next in self.link[min_node]:
                if next not in self.S:
                    if next not in pre or min_edge_len + self.link[min_node][next] < self.S[pre[next]] + self.link[pre[next]][next]:
                        self.U_que.put_nowait((min_edge_len + self.link[min_node][next], next))
                        pre[next] = min_node

        # 生成路径函数
        def __getPath(end_node):
            stack = []
            node = end_node
            while True:
                stack.append(node)
                if node not in pre:
                    break
                node = pre[node]
            return stack[::-1]

        return [(k, v, __getPath(k)) for k, v in self.S.items()] if get_path else [(k, v) for k, v in self.S.items()]


node_num, edge_num = map(int, input().split())


link = {}
rev_edges = []
rev = {}
for _ in range(edge_num):
    a, b, w = map(int, input().split())
    if a not in link:
        link[a] = []
    link[a].append((b, w))

    if (b, a) not in rev:
        rev[(b, a)] = w
    elif w < rev[(b, a)]:
        rev[(b, a)] = w


for (s, e), w in rev.items():
    rev_edges.append((s, e, w))

start, end, k = map(int, input().split())


# 验证起点到终点的连通性
def dfs(cur, end, visited) -> bool:
    if cur == end:
        return True

    if cur in link:
        for next_node, _ in link[cur]:
            if next_node not in visited:
                visited.add(next_node)
                if dfs(next_node, end, visited):
                    return True

    return False



if not dfs(start, end, set()):
    print(-1)
else:
    min_dist = DijkStra(end, rev_edges, [i for i in range(1, node_num+1)], is_directed=True)
    min_dist_list = min_dist.getMinPathLen()
    min_dis_map = {node: dis for node, dis in min_dist_list}
    #print(min_dis_map)


    from queue import PriorityQueue
    min_heap = PriorityQueue()

    min_heap.put( (min_dis_map[start], 0, start) )
    cnt = 0
    while not min_heap.empty():
        _, cur_dis, cur_node = min_heap.get()
        #print(cur_node, cur_dis)
        if cur_node == end and cur_dis != 0:
            cnt += 1
            if cnt == k:
                print(cur_dis)
                break

        if cur_node not in link:
            continue

        for next_node, edge_w in link[cur_node]:
            min_heap.put((cur_dis + edge_w + min_dis_map[next_node], cur_dis + edge_w, next_node))

    if cnt != k:
        print(-1)