第K短路(Dijkstra+A*)

给定一张 N 个点（编号 1,2…N），M 条边的有向图，求从起点 S 到终点 T 的第 K 短路的长度，路径允许重复经过点或边。
注意： 每条最短路中至少要包含一条边。

输入格式
第一行包含两个整数 N 和 M。
接下来 M 行，每行包含三个整数 A,B 和 L，表示点 A 与点 B 之间存在有向边，且边长为 L。
最后一行包含三个整数 S,T 和 K，分别表示起点 S，终点 T 和第 K 短路。
输出格式
输出占一行，包含一个整数，表示第 K 短路的长度，如果第 K 短路不存在，则输出 −1。

数据范围
1≤S,T≤N≤1000,
0≤M≤104,
1≤K≤1000,
1≤L≤100

输入样例：
2 2
1 2 5
2 1 4
1 2 2
输出样例：
14

#include<cstring>
#include<algorithm>
#include<iostream>
#include<queue>
using namespace std;

typedef pair<int,int>   PII;
typedef pair<int,PII>   PIII;

const int N=1010,M=200010;

int n,m;
int h[N],rh[N],w[M],e[M],ne[M],idx;
int dist[N],f[N],vis[N];
int S,T,K;
void add(int *h,int a,int b,int c){
    e[idx]=b;w[idx]=c;ne[idx]=h[a];h[a]=idx++;
}
void dijkstra(){
    priority_queue<PII,vector<PII>,greater<PII>>heap;
    memset(dist,0x3f,sizeof dist);
    dist[T]=0;
    heap.push({0,T});
    
    while(heap.size()){
        auto t=heap.top();
        heap.pop();
        int vertex=t.second;
        if(vis[vertex]) continue;
        vis[vertex]=1;
        
        for(int i=rh[vertex];~i;i=ne[i]){
            int j=e[i];
            if(dist[j]>dist[vertex]+w[i]){
                dist[j]=dist[vertex]+w[i];
                heap.push({dist[j],j});
            }
        }
    
        
    }
    memcpy(f,dist,sizeof dist);
}
int A_star(){
    priority_queue<PIII,vector<PIII>,greater<PIII>>heap;
    heap.push({f[S],{0,S}});
    memset(vis,0,sizeof vis);
    
    while(heap.size()){
        auto t=heap.top();
        heap.pop();
        int vertex=t.second.second,realdist=t.second.first;
        if(vis[vertex]>=K)  continue;
        vis[vertex]++;
        if(vertex==T&&vis[vertex]==K)   return realdist;
        
        for(int i=h[vertex];~i;i=ne[i]){
            int j=e[i];
            if(vis[j]<K)    
                heap.push({realdist+w[i]+f[j],{realdist+w[i],j}});
        }
        
    }
    
    return -1;
}
int main(){
    scanf("%d%d",&n,&m);
    memset(h,-1,sizeof h);
    memset(rh,-1,sizeof rh);
    while(m--){
        int a,b,c;
        scanf("%d%d%d",&a,&b,&c);
        add(h,a,b,c);add(rh,b,a,c);
    }
    scanf("%d%d%d",&S,&T,&K);
    if(S==T)    K++;
    dijkstra();
    printf("%d\n",A_star());
}