CodeForces - 337D 树形dp

题意：一颗树上有且仅有一只恶魔，恶魔会污染距离它小于等于d的点，现在已经知道被污染的m个点，问恶魔在的可能结点的数量。

容易想到，要是一个点到（距离最远的两个点）的距离都小于等于d，那么这个点就有可能是恶魔所在的点（虽然我也是没有证明凭感觉，逃~~）。

那么问题就难在怎么快速找到这m个点中距离最远的两个点？我们会想到这跟 找树的直径非常相似，于是就是用找树的直径的方法：树形dp找出这两个点。

然后距离这两个点的距离都小于等于d的就统计答案即可。

然后要注意m=1的情况。

细节请看代码：

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+;
int n,m,d,num,ans,rec,v[N];
vector<int> G[N];

int d1[N],d2[N],r1[N],r2[N];
void dfs1(int x,int fa) {
    if (v[x]) d1[x]=,r1[x]=x;
    for (int i=;i<G[x].size();i++) {
        int y=G[x][i];
        if (y==fa) continue;
        dfs1(y,x);
        if (d1[y]+>=d1[x]) {
            d2[x]=d1[x]; r2[x]=r1[x];
            d1[x]=d1[y]+; r1[x]=r1[y];
        }
        else if (d1[y]+>d2[x]) {
            d2[x]=d1[y]+; r2[x]=r1[y];
        }
    }
    if (d1[x]+d2[x]>ans) {
        ans=d1[x]+d2[x];
        rec=x;
    }
}

int dep1[N],dep2[N];
void dfs2(int x,int dd,int fa,int *dep) {
    dep[x]=dd;
    for (int i=;i<G[x].size();i++) {
        int y=G[x][i];
        if (y==fa) continue;
        dfs2(y,dd+,x,dep);
    }
}

int main()
{
    cin>>n>>m>>d;
    for (int i=;i<=m;i++) {
        int x; scanf("%d",&x);
        v[x]=; num=x;
    }
    for (int i=;i<n;i++) {
        int x,y; scanf("%d%d",&x,&y);
        G[x].push_back(y);
        G[y].push_back(x);
    }

    memset(d1,-0x3f,sizeof(d1));
    memset(d2,-0x3f,sizeof(d2));
    ans=; rec=;
    dfs1(,);

    int cnt=;
    if (m>) {
        dfs2(r1[rec],,,dep1);
        dfs2(r2[rec],,,dep2);
        for (int i=;i<=n;i++)
            if (dep1[i]<=d && dep2[i]<=d) cnt++;
    } else {
        dfs2(num,,,dep1);
        for (int i=;i<=n;i++)
            if (dep1[i]<=d) cnt++;
    }
    cout<<cnt<<endl;
    return ;
}