A*求单元第k短路的模板

#include <iostream>
#include <cstdio>
#include <fstream>
#include <algorithm>
#include <cmath>
#include <deque>
#include <vector>
#include <list>
#include <queue>
#include <string>
#include <cstring>
#include <map>
#define PI acos(-1.0)
#define mem(a,b) memset(a,b,sizeof(a))
#define sca(a) scanf("%d",&a)
#define pri(a) printf("%d\n",a)
#define MM 500002
#define MN 1002
#define INF 168430090
using namespace std;
typedef long long ll;
int n,m,st,en,k,cnt,tmp,vis[MN],d[MN],q[MM*5],head[MN],rhead[MN];
struct edge
{
int v,w,next;
}e[MM],re[MM];//一个正向，一个反向
struct A
{
int f,g,v;
bool operator < (const A a)const
{
if(a.f==f) return a.g<g;
return a.f<f;
}
};
void add(int u,int v,int w)
{
e[cnt].v=v; e[cnt].w=w; //正向邻接表
e[cnt].next=head[u]; head[u]=cnt;
re[cnt].v=u; re[cnt].w=w; //反向邻接表，为求的是h(p)函数
re[cnt].next=rhead[v]; rhead[v]=cnt++;
}
void spfa()
{
int i,l=0,r=1,u,v,w;
for(i=1;i<=n;i++) d[i]=INF;
mem(vis,0);
q[0]=en; d[en]=0;
while(l<r) //模拟队列，如果用队列可能会RE
{
u=q[l++]; vis[u]=0;
for(i=rhead[u];i!=-1;i=re[i].next)
{
v=re[i].v; w=re[i].w;
if(d[v]>d[u]+w)
{
d[v]=d[u]+w; //更新每点的最短路，h(p)在本题中就是d数组
if(!vis[v])
{
q[r++]=v;
vis[v]=1;
}
}
}
}
}
int Astar()
{
priority_queue<A>Q;
if(st==en) k++; //这句WA了一发，因为源点与终点相同时少了本身的一次，所以要加上
if(d[st]==INF) return -1;
A s1,s2;
s1.v=st; s1.g=0; s1.f=s1.g+d[st]; //即f(p)=g(p)+h(p)
Q.push(s1);
while(!Q.empty())
{
s2=Q.top(); Q.pop();
if(s2.v==en)
{
tmp++;
if(tmp==k) return s2.g;
}
for(int i=head[s2.v];i!=-1;i=e[i].next)
{
s1.v=e[i].v; s1.g=s2.g+e[i].w; s1.f=s1.g+d[s1.v];
Q.push(s1);
}
}
return -1;
}
int main()
{
int i,u,v,w;
scanf("%d%d",&n,&m);
mem(head,-1); mem(rhead,-1);
for(i=0;i<m;i++)
{
scanf("%d%d%d",&u,&v,&w);
add(u,v,w);
}
scanf("%d%d%d",&st,&en,&k);
spfa();
printf("%d\n",Astar());
return 0;
}