UVA10917 Walk Through the Forest (dijkstra + dfs)

#include <cstdio>
#include <queue>
#include <cstring>
#include <algorithm>
#define LL long long
#define inf 1LL << 60
using namespace std;

const int maxn = 1e6+ 10;
int N,M,cnt;
int head[maxn],vis[1100],dp[1100];
LL dist[maxn];
struct Edge{
int to,w,next;
}edge[maxn];
struct node{
int dis,v;
friend bool operator <(node A, node B){
return A.dis > B.dis;
}
};

void add(int u, int v, int w){
edge[cnt].to = v;
edge[cnt].w = w;
edge[cnt].next = head[u];
head[u] = cnt ++;
}

void dijkstra(int x){
for(int i = 1; i <= N; i ++ ) dist[i] = inf,vis[i] = 0;
dist[x] = 0;
priority_queue<node > q;
node p;
p.dis = 0, p.v = x;
q.push(p);
while(!q.empty()){
p = q.top(); q.pop();
if(vis[p.v]) continue;
vis[p.v] = 1;
for(int i = head[p.v]; i != -1; i = edge[i].next){
int v = edge[i].to;
if(dist[v] > dist[p.v] + edge[i].w){
dist[v] = dist[p.v] + edge[i].w;
node no;
no.dis = dist[v], no.v = v;
q.push(no);
}
}
}
}

int dfs(int x){
if(x == 2) return 1;
if(dp[x]) return dp[x];
int ret = 0;
for(int i = head[x]; i != -1; i = edge[i].next){
int v = edge[i].to;
if(dist[x] > dist[v]) ret += dfs(v);
}
return dp[x] = ret;
}

int main(){
while(~scanf("%d",&N) && N){
scanf("%d",&M);
memset(head,-1,sizeof(head));
int u,v,w;
cnt = 0;
for(int i = 1; i <= M; i ++){
scanf("%d%d%d",&u,&v,&w);
add(u,v,w);
add(v,u,w);
}
dijkstra(2);
memset(dp,0,sizeof(dp));
printf("%d\n",dfs(1));
}
return 0;
}