ACWing：178. 第K短路 （A*算法）

178. 第K短路 - AcWing题库

ac代码：

#include<iostream>
#include<cstring>
#include<queue>
using namespace std;
const int N=1010;
const int M=20020;
struct node{
    int d,end,d1;
    bool operator <(const node &x)const{
        return d>x.d;
    }
};
priority_queue<node>pq;
int dis[N],vis[N],h1[M],h2[M],ne[M],v[M],w[M],cnt[N];
int n,m,s,t,k,idx;
int ct=0;
void add(int h[],int ui,int vi,int wi){
    ne[++idx]=h[ui]; h[ui]=idx;
    v[idx]=vi;w[idx]=wi;
}
void dij(){
    pq.push({0,t});
    while(pq.size()){
        node tmp=pq.top(); pq.pop();
        if(vis[tmp.end]) continue;
        vis[tmp.end]=1; dis[tmp.end]=tmp.d;
        for(int i=h2[tmp.end];i;i=ne[i]){
            if(vis[v[i]]) continue;
            pq.push({tmp.d+w[i],v[i]});
        }
    }
} 
int A(){
    pq.push({dis[s],s,0});
    while(pq.size()){
        node tmp=pq.top(); pq.pop();
        cnt[tmp.end]++;
        if(tmp.end==t){  ct++; if(ct==k) return tmp.d;}
        if(cnt[tmp.end]>k) continue;
        for(int i=h1[tmp.end];i;i=ne[i]){
            pq.push({tmp.d1+w[i]+dis[v[i]],v[i],tmp.d1+w[i]});
        }
    }
    return -1;
}
int main(){
    scanf("%d%d",&n,&m);
    while(m--){
        int ui,vi,wi;
        scanf("%d%d%d",&ui,&vi,&wi);
        add(h1,ui,vi,wi);
        add(h2,vi,ui,wi);
    }
    scanf("%d%d%d",&s,&t,&k);
    if(s==t) k++;
    dij(); 
    printf("%d",A());
    return 0;
}