AcWing178 第K短路（A*）

#include<iostream>
#include<algorithm>
#include<string>
#include<cstring>
#include<cstdio>
#include<queue>
using namespace std;
typedef pair<int,int> pll;
typedef pair<int,pair<int,int>> plll;
int n,m,idx;
int f[10005],rf[10005];
int h[1010];
int e[200001],ne[200001],w[200001];
int dis[1010];
int st[10005];
int s,T,k;
void add(int *h,int a,int b,int l){
    e[idx]=b,w[idx]=l,ne[idx]=h[a],h[a]=idx++;
}
void dj(int s){
    priority_queue<pll,vector<pll>,greater<pll>> q;
    memset(dis, 0x3f, sizeof dis);
    dis[s]=0;
    q.push({dis[s],s});
    while(!q.empty()){
        auto t=q.top();
        q.pop();
        int i,j;
        if(st[t.second])
        continue;
        st[t.second]=1;
        for(i=rf[t.second];~i;i=ne[i]){
            j=e[i];
            if(dis[j]>w[i]+dis[t.second]){
                dis[j]=w[i]+dis[t.second];
                q.push({dis[j],j});
            }
        }
    }
    memcpy(h, dis, sizeof dis);
}
int astar(int s){
    priority_queue<plll,vector<plll>,greater<plll>> q;
    int distance=0;
    q.push({h[s],{0,s}});
    memset(st,0,sizeof(st));
    while(q.size()){
        auto t=q.top();
        q.pop();
        int ver=t.second.second;
        distance=t.second.first;
        st[ver]++;
        if (ver==T&&st[ver]==k) 
        return distance;

        int i,j;
        for(i=f[ver];~i;i=ne[i]){
            j=e[i];
            if(st[j]<k)
            q.push({h[j]+w[i]+distance,{distance+w[i],j}});
        }
        
    }
    return -1;
    
    
}
int main(){
    int a,b,l;
    cin>>n>>m;
    int i,j;
    memset(f, -1, sizeof f);
    memset(rf, -1, sizeof rf);
    for(i=0;i<m;i++){
        cin>>a>>b>>l;
        add(f,a,b,l);
        add(rf,b,a,l);
    }
    cin>>s>>T>>k;
    if(s==T)
    k++;
    dj(T);
    cout<<astar(s)<<endl;
}