codeforces 337D D. Book of Evil (树形 dp)

题解 ： 我们发现可能藏着魔鬼的点就是那些，到所有被感染的点的最大距离小于等于d的点，剩下的我们的目的就是求出每个点到所有感染点的最大距离，这就是一个类似于求，树上的任意一个点，到所有点的最大距离，就是一个简单的树形dp;

#include <algorithm>
#include <cstring>
#include <cstdio>
#include <vector>
#include <iostream>

using namespace std;
const int maxn = 1e6 + 10;
const int INF =  1e9 + 7;
int n,m,d;
vector <int> tr[maxn];
vector <int> G[maxn];
bool tag[maxn] = {0};
int f[maxn] = {0};
int g[maxn] = {0};
bool vis[maxn] = {0};
int dp[maxn] = {0};
int pre[maxn] = {0};
int temp[maxn] = {0};

void build (int root) {
    vis[root] = 1;
    int cnt = tr[root].size();
    for (int i = 0;i < cnt; ++ i) {
        int u = tr[root][i];
        if (!vis[u]) {
            pre[u] = root;
            G[root].push_back(u);
            build (u);
        }
    }
}

int dfs1 (int root) {
    if (tag[root]) f[root] = 0;
    else f[root] = -INF;
    int cnt = G[root].size ();
    if (cnt == 0) {
        return f[root];
    }
    for (int i = 0;i < cnt; ++ i) {
        int u = G[root][i];
        dfs1 (u);
        f[root] = max (f[root],f[u] + 1);
    }
    return f[root];
}

void dfs2 (int root) {
    g[root] = g[pre[root]];
    if (f[root] + 1 != temp[pre[root]])
    g[root] = max (g[root],temp[pre[root]]);
    else {
        int cnt = G[pre[root]].size();
        for (int i = 0;i < cnt; ++ i) {
            int u = G[pre[root]][i];
            if (u == root) continue;
            g[root] = max (g[root],f[u] + 1);
        }
    }
    g[root] = g[root] + 1;
    if (tag[root]) g[root] = max (0,g[root]);
    int cnt = G[root].size();
    for (int i = 0;i < cnt; ++ i) {
        int u = G[root][i];
        dfs2 (u);
    }
}

int main () {
    scanf ("%d%d%d",&n,&m,&d);
    for (int i = 1;i <= m; ++ i) {
        int x;
        scanf ("%d",&x);
        tag[x] = 1;
    }
    for (int i = 1;i < n; ++ i) {
        int x,y;
        scanf ("%d%d",&x,&y);
        tr[x].push_back(y);
        tr[y].push_back(x);
    }
    for (int i = 0;i <= n; ++ i) temp[i] = -INF;
    build(1);
    dfs1 (1);
    for (int i = 1;i <= n; ++ i) {
        int cnt = G[i].size();
        if (cnt) {
            for (int j = 0;j < cnt; ++ j) {
                int u = G[i][j];
                temp[i] = max (temp[i],f[u] + 1);
            }
        }
    }
    g[0] = -INF;
    dfs2 (1);
    for (int i = 1;i <= n; ++ i) dp[i] = max (f[i],g[i]);
    int ans = 0;
    for (int i = 1;i <= n; ++ i) {
        if (dp[i] >= 0 && dp[i] <= d) {
            ans ++;
        }
    }
    printf ("%d\n",ans);
    return 0;
}