POJ3463-dij求次短路径

/*
抄也罢，袭也好，终究还是过了
两种方法：
1.A*求k短，k=2；(实际上爆内存)
2.延拓dijkstra求最短到次短(最终选用的方法，32ms)
*/
#include <cstdio>
#include <iostream>
#include <cstring>
#include <queue>
#define mem(a,x) memset(a,x,sizeof(a))
using namespace std;

const int NN=1005;
const int MM=10005;
const int INF=0x3fffffff;

struct Edge{
   int v,dis,next;
}edge[MM];
int head[NN],ecnt,n,m,S,F,k;
int dis[NN][2],cou[NN][2];  //dis[NN][0]记录最短，dis[NN][1]记录次短，cou记录相应的路径条数
struct node{
    int v,flag;
    node(int vv,int f){v=vv,flag=f;}
    bool operator<(const node &p)const {
        return dis[v][flag]>dis[p.v][p.flag];
    }
};

void addedge(int u,int v,int dis)
{
    edge[ecnt].v=v;
    edge[ecnt].dis=dis;
    edge[ecnt].next=head[u];
    head[u]=ecnt++;
}

void init()
{
    ecnt=0;
    mem(head,-1);
    scanf("%d%d",&n,&m);
    int x,y,z;
    for (int i=1; i<=m; i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        addedge(x,y,z);
    }
    scanf("%d%d",&S,&F);
}

bool vis[NN][2];
int dij()
{
    mem(vis,0);
    mem(cou,0);
    for (int i=1; i<=n; i++) dis[i][0]=dis[i][1]=INF;
    priority_queue<node> q;
    q.push(node(S,0));
    dis[S][0]=0; cou[S][0]=1;
    while (!q.empty())
    {
        int u=q.top().v,flag=q.top().flag;  //flag标记是最短还是次短
        q.pop();
        if (vis[u][flag]) continue;   //vis过的最短或次短已经是最优解了，不用再作为更新的点了
        vis[u][flag]=true;
        for (int i=head[u]; i!=-1; i=edge[i].next)
        {
            int v=edge[i].v;
            int w=dis[u][flag]+edge[i].dis;
            if (w<dis[v][0])
            {
                if (dis[v][0]!=INF)
                {
                    dis[v][1]=dis[v][0];
                    cou[v][1]=cou[v][0];
                    q.push(node(v,1));
                }
                dis[v][0]=w;
                cou[v][0]=cou[u][flag];
                q.push(node(v,0));
            }
            else if (w==dis[v][0]) cou[v][0]+=cou[u][flag];
            else if (w<dis[v][1])
            {
                dis[v][1]=w;
                cou[v][1]=cou[u][flag];
                q.push(node(v,1));
            }
            else if (w==dis[v][1]) cou[v][1]+=cou[u][flag];
        }
    }
    int ret=cou[F][0];
    if (dis[F][1]==dis[F][0]+1) ret+=cou[F][1];
    return ret;
}

int main()
{
    int cas;
    scanf("%d",&cas);
    while (cas--)
    {
        init();
        printf("%d\n",dij());
    }
    return 0;
}