A* 第k短路

#include<iostream>
#include<cstring>
#include<queue>

using namespace std;

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

const int N=1010,M=2e5+10;

int ne[M],e[M],w[M],rh[N],h[N],idx;
int dist[N],cnt[N];
int n,m,S,T,K;
bool st[N];

void add(int h[],int x,int y,int z){
    e[idx]=y;
    w[idx]=z;
    ne[idx]=h[x];
    h[x]=idx++;
}

void dijkstra(){

    memset(dist,0x3f,sizeof dist);
    dist[T]=0;

    priority_queue<PII,vector<PII>,greater<PII>> heap;
    heap.push({0,T});

    while(heap.size()){
        auto k=heap.top();
        heap.pop();

        int ver=k.second,distance=k.first;

        if(st[ver])continue;
        st[ver]=true;

        for(int i=rh[ver];i!=-1;i=ne[i]){
            int j=e[i];
            if(dist[j]>distance+w[i]){
                dist[j]=distance+w[i];
                heap.push({dist[j],j});
            }
        }
    }
}

int astar(){
    priority_queue<PIII,vector<PIII>,greater<PIII>> heap;
    heap.push({dist[S],{0,S}});

    while(heap.size()){
        auto k=heap.top();
        heap.pop();

        int distance=k.second.first,ver=k.second.second;

        cnt[ver]++;
        if(cnt[T]==K)return distance;

        for(int i=h[ver];i!=-1;i=ne[i]){
            int j=e[i];
            if(cnt[j]<K){
                heap.push({distance+w[i]+dist[j],{distance+w[i],j}});
            }
        }
    }
    return -1;
}

int main(){

    cin>>n>>m;

    memset(h,-1,sizeof h);
    memset(rh,-1,sizeof rh);

    while(m--){
        int x,y,z;
        cin>>x>>y>>z;

        add(h,x,y,z);
        add(rh,y,x,z);

    }

    cin>>S>>T>>K;

    if(S==T)K++;

    dijkstra();

    cout<<astar();

    return 0;
}