Codeforce 337D(树的直径变形)

链接：点击打开链接

题意：在一个n个点的树中，给出m个点，问这颗树中距离这m个点的距离小于等于d的点的个数

代码：

#include <vector>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <iostream>
#include <algorithm>
using namespace std;
const int INF=0x3f3f3f3f;
struct node{
    int to,cost;
    node(int x,int y){
        to=x,cost=y;
    }
};
vector<node> G[100005];
int n,dp[100005],vis[100005],dis[100005][5],sign[100005];
void dfs1(int s){
    int i,tmp;
    sign[s]=1;
    for(i=0;i<G[s].size();i++){
        tmp=G[s][i].to;
        if(!sign[tmp]){
        dfs1(tmp);
        if(dis[tmp][0]+G[s][i].cost>dis[s][0]){
            dis[s][0]=dis[tmp][0]+G[s][i].cost;
            vis[s]=tmp;
        }
        }
    }
    for(i=0;i<G[s].size();i++){
        tmp=G[s][i].to;
        if(tmp!=vis[s])
        dis[s][1]=max(dis[s][1],dis[tmp][0]+G[s][i].cost);
    }
}
void dfs2(int s){
    int i,tmp;
    sign[s]=1;
    for(i=0;i<G[s].size();i++){
        tmp=G[s][i].to;
        if(!sign[tmp]){
        if(tmp==vis[s])
        dp[tmp]=max(dp[s],dis[s][1])+G[s][i].cost;
        else
        dp[tmp]=max(dp[s],dis[s][0])+G[s][i].cost;
        dfs2(tmp);
        }
    }
}
int main(){                                     //就是树的直径变形，有原来向上和向下
    int i,j,x,y,m,d,ans,num;                    //的最长长度变为求这个点到给定的m个点中
    while(scanf("%d%d%d",&n,&m,&d)!=EOF){       //的最长向上和向下的距离
        memset(vis,0,sizeof(vis));
        memset(sign,0,sizeof(sign));            //可以先看树的直径
        for(i=1;i<=n;i++){                      //http://blog.csdn.net/stay_accept/article/details/50813783
        G[i].clear();
        dp[i]=dis[i][0]=dis[i][1]=-INF;
        }
        for(i=1;i<=m;i++){
            scanf("%d",&num);
            dp[num]=dis[num][0]=dis[num][1]=0;
        }
        for(i=2;i<=n;i++){
            scanf("%d%d",&x,&y);
            G[x].push_back(node(y,1));
            G[y].push_back(node(x,1));
        }
        dfs1(1);
        memset(sign,0,sizeof(sign));
        dfs2(1);
        ans=0;
        for(i=1;i<=n;i++)
        if(dis[i][0]<=d&&dp[i]<=d)
        ans++;
        printf("%d\n",ans);
    }
    return 0;
}