[CF337D] Book of Evil

[CF337D] Book of Evil - 树形dp

Description

有一棵树有 (n) 个节点，其中有 (m) 个节点发现了怪物。已知树上有一本魔法书，魔法书可以让到其距离小于等于 (d) 的点出现怪物，求魔法书所在点有几种可能。

Solution

对每个点，我们要求出三个值：子树内最远怪物，与最远不在同一个分支内的次远怪物（定义有点难以描述），子树外最远怪物（的距离）

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 1000005;

struct Node
{
    int max = INT_MIN;
    int sec = INT_MIN;
    int out = INT_MIN;
} node[N];

int n, m, d, a[N];

vector<int> g[N];

void dfs(int p, int from)
{
    if (a[p])
    {
        node[p].max = node[p].sec = 0;
    }
    for (int q : g[p])
        if (q != from)
        {
            dfs(q, p);
            if (node[p].max < node[q].max + 1)
            {
                node[p].sec = node[p].max;
                node[p].max = node[q].max + 1;
            }
            else
            {
                node[p].sec = max(node[p].sec, node[q].max + 1);
            }
        }
}

void dfsout(int p, int from)
{
    for (int q : g[p])
        if (q != from)
        {
            if (node[p].max == node[q].max + 1)
            {
                node[q].out = max(node[p].out + 1, node[p].sec + 1);
            }
            else
            {
                node[q].out = max(node[p].out + 1, node[p].max + 1);
            }
            dfsout(q, p);
        }
}

signed main()
{
    ios::sync_with_stdio(false);

    cin >> n >> m >> d;
    for (int i = 1; i <= m; i++)
    {
        int x;
        cin >> x;
        a[x] = 1;
    }

    for (int i = 1; i < n; i++)
    {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    dfs(1, 0);
    dfsout(1, 0);

    int ans = 0;
    for (int i = 1; i <= n; i++)
    {
        int p = i;
        if (i == 1)
        {
            ans += node[p].max <= d && node[p].sec <= d;
        }
        else if (node[p].max <= d && node[p].out <= d)
        {
            ++ans;
        }
    }

    cout << ans << endl;
}