CodeForces - 337D Book of Evil （树形dp）

题意：有一颗有n个结点的树，树上存在一个污染源（位置不确定），它可以污染与它距离不超过d的节点，现给出m个被污染的节点（污染源本身也可能是被污染的节点），求污染源可能的位置数。

思路：树形dp，树的最长链，dp[i][0]表示i到i的子树中最远污染源的距离，dp[i][1]表示i到i的子树中次远污染源的距离，dp[i][2]表示i经过i的父节点到最远污染源的距离。那么污染源到这些被污染的点的最远距离应该不大于d,枚举一遍统计合法的点即为答案。

其他细节看代码。

#include<bits/stdc++.h>
using namespace std;
int dp[100005][3];//dp[i][0]表示i到i的子树中最远污染源的距离，dp[i][1]表示i到i的子树中次远污染源的距离，dp[i][2]表示i经过i的父节点到最远污染源的距离.
bool vis[100005];
vector<int>vec[100005];
void dfs1(int x,int y) {
    int k=vec[x].size();
    for(int i=0; i<k; i++) {
        int v=vec[x][i];
        if(v==y)
            continue;
        if(vis[v])//该点被污染更新dp[v][0]
            dp[v][0]=0;
        dfs1(v,x);//先更新子树中的答案
        //dp[v][0]>=0说明v的子树中存在被污染的点，更新父结点的答案
        if(dp[v][0]>=0) {
            if(dp[v][0]+1>=dp[x][0]) {
                dp[x][1]=dp[x][0];
                dp[x][0]=dp[v][0]+1;
            } else if(dp[v][0]+1>dp[x][1])
                dp[x][1]=dp[v][0]+1;
        }
    }
}
void dfs2(int x,int y) {
    int k=vec[x].size();
    //x为被污染点，应更新答案，但存在x的父结点中有被污染点已经更新答案，所以取最大的答案
    if(vis[x])
        dp[x][2]=max(0,dp[x][2]);
    for(int i=0; i<k; i++) {
        int v=vec[x][i];
        if(v==y)
            continue;
        int ans=max(dp[x][2],(dp[x][0]==dp[v][0]+1)?dp[x][1]:dp[x][0]);
        //答案合法，更新子节点的答案
        if(ans>=0)
            dp[v][2]=ans+1;
        dfs2(v,x);
    }
}
int main() {
    int n,m,d;
    scanf("%d%d%d",&n,&m,&d);
    for(int i=0; i<m; i++) {
        int a;
        scanf("%d",&a);
        vis[a]=true;
    }
    for(int i=0; i<n-1; i++) {
        int a,b;
        scanf("%d%d",&a,&b);
        vec[a].push_back(b);
        vec[b].push_back(a);
    }
    for(int i=1; i<=n; i++) {
        dp[i][0]=dp[i][1]=dp[i][2]=-10;
    }
    dfs1(1,0);
    dfs2(1,0);
    int sum=0;
    //cout<<dp[2][2]<<endl;
    for(int i=1; i<=n; i++) {
        if(max(dp[i][0],dp[i][2])<=d)
            sum++;
    }
    printf("%d\n",sum);
    return 0;
}