Codeforces 337D 树形dp + 小技巧

题意：

题目链接：http://codeforces.com/problemset/problem/337/D
一棵n个节点的树，树上有m个点有恶魔，现在在树上要定位一个魔法书的位置，已知每个恶魔离魔法书的距离不超过d，问一共有多少个节点可能有魔法书。

---

思路：

典型的树形dp，需要求出每个节点到其他恶魔点的最远距离dis，最后看有多少个点可以满足dis <= d即可。
分两次dfs，一次从下到上，一次从上到下。
dp1[u]表示节点u到它自己这棵子树中恶魔点的最远距离。
dp2[u]表示节点u到不在它自己子树的恶魔点的最远距离。
需要注意的地方是，在第二次dfs中，遍历到节点u，需要求出节点u的兄弟之中dp1的最大值，如果暴力枚举每个兄弟会T。这里有个小技巧，实现保存下来每个点的儿子中dp1最大值的编号id，在dfs2中，如果u==id[pa]再暴力枚举，否则直接取父亲pa的最大值即可。

---

代码：

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 1e5 + 10;
const int INF = 0x3f3f3f3f;

vector <int> tree[MAXN];
int dp1[MAXN], dp2[MAXN], id[MAXN];
bool evil[MAXN];

void dfs1(int u, int pa) {
    int cnt = tree[u].size();
    int res = evil[u] ? 0 : -1, Max = -INF;
    for (int i = 0; i < cnt; i++) {
        int v = tree[u][i];
        if (v == pa) continue;
        dfs1(v, u);
        if (dp1[v] == -1) continue;
        res = max(dp1[v] + 1, res);
        if (Max < dp1[v]) {
            Max = dp1[v];
            id[u] = v;
        }
    }
    dp1[u] = res;
}

void dfs2(int u, int pa, int fa) {
    int res = evil[u] ? 0 : -1;
    if (pa != -1) {
        if (dp2[pa] != -1)
            res = max(res, dp2[pa] + 1);
        if (dp1[pa] != -1) {
            if (u == id[pa]) {
                int cnt = tree[pa].size();
                for (int i = 0; i < cnt; i++) {
                    int bro = tree[pa][i];
                    if (bro == u || bro == fa) continue;
                    if (dp1[bro] == -1) continue;
                    res = max(res, dp1[bro] + 2);
                }
            }
            else {
                res = max(res, dp1[pa] + 1);
            }
        }
    }
    dp2[u] = res;
    int cnt = tree[u].size();
    for (int i = 0; i < cnt; i++) {
        int v = tree[u][i];
        if (v == pa) continue;
        dfs2(v, u, pa);
    }
}

int main() {
    //freopen("in.txt", "r", stdin);
    int n, m, d;
    scanf("%d%d%d", &n, &m, &d);
    for (int i = 1; i <= m; i++) {
        int x;
        scanf("%d", &x);
        evil[x] = true;
    }
    for (int i = 1; i <= n; i++) tree[i].clear();
    for (int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    dfs1(1, -1);
    dfs2(1, -1, -1);
    int ans = 0;
    for (int i = 1; i <= n; i++) {
        if (max(dp1[i], dp2[i]) <= d)
            ++ans;
    }
    printf("%d\n", ans);
    return 0;
}