CodeForces 337D :Book of Evil 树形DP

传送门

题目描述

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

分析

两种做法

• 树形DP
  我们可以两次dfs预处理出来每个点离怪物最远的距离
  $f[u][0]$ :$u$到子树中怪物的最远距离
  $f[u][1]$ :$u$到子树中怪物的次远距离
  $f[u][2]$ :$u$到父节点中怪物的最远距离
  然后判断每个点的距离就可以了

• 树的直径
  先空着有时间再补

代码

#pragma GCC optimize(3)
#include <bits/stdc++.h>
#define debug(x) cout<<#x<<":"<<x<<endl;
#define dl(x) printf("%lld\n",x);
#define di(x) printf("%d\n",x);
#define _CRT_SECURE_NO_WARNINGS
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(),(x).end()
#define fi first
#define se second
#define SZ(x) ((int)(x).size())
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> PII;
typedef vector<int> VI;
const int INF = 0x3f3f3f3f;
const int N = 2e6 + 10;
const ll mod= 1000000007;
const double eps = 1e-9;
const double PI = acos(-1);
template<typename T>inline void read(T &a){char c=getchar();T x=0,f=1;while(!isdigit(c)){if(c=='-')f=-1;c=getchar();}
while(isdigit(c)){x=(x<<1)+(x<<3)+c-'0';c=getchar();}a=f*x;}
int gcd(int a,int b){return (b>0)?gcd(b,a%b):a;}
int h[N],ne[N],e[N],idx;
int f[N][3];
int n,m,d;
bool st[N];

void add(int x,int y){
    ne[idx] = h[x],e[idx] = y,h[x] = idx++;
}

void dfs1(int u,int fa){
    for(int i = h[u];~i;i = ne[i]){
        int v = e[i];
        if(v == fa) continue;
        dfs1(v,u);
        if(f[v][0] >= 0){
            if(f[v][0] + 1 >= f[u][0]){
                f[u][1] = f[u][0];
                f[u][0] = f[v][0] + 1;
            }
            else if(f[v][0] + 1 > f[u][1])
                f[u][1] = f[v][0] + 1;
        } 
    }
}

void dfs2(int u,int fa){
    for(int i = h[u];~i;i = ne[i]){
        int v = e[i];
        if(v == fa) continue;
        int t = 0;
        if(f[u][0] == f[v][0] + 1) t = max(f[u][1],f[u][2]);
        else t = max(f[u][0],f[u][2]);
        if(t >= 0) f[v][2] = t + 1;
        dfs2(v,u);
    }
}

int main(){
    memset(h,-1,sizeof h);
    read(n),read(m),read(d);
    memset(f,-0x3f,sizeof f);
    for(int i = 1;i <= m;i++){
        int x;
        read(x);
        st[x] = 1;
        f[x][0] = 0;
        f[x][2] = 0;
    }
    for(int i = 1;i < n;i++){
        int x,y;
        read(x),read(y);
        add(x,y),add(y,x);
    }
    dfs1(1,-1);
    dfs2(1,-1);
    int ans = 0;
    for(int i = 1;i <= n;i++) if(max(f[i][0],f[i][2]) <= d) ans++;
    di(ans);
    return 0;
}

/**
*　　┏┓　　　┏┓+ +
*　┏┛┻━━━┛┻┓ + +
*　┃　　　　　　　┃
*　┃　　　━　　　┃ ++ + + +
*  ████━████+
*  ◥██◤　◥██◤ +
*　┃　　　┻　　　┃
*　┃　　　　　　　┃ + +
*　┗━┓　　　┏━┛
*　　　┃　　　┃ + + + +Code is far away from 　
*　　　┃　　　┃ + bug with the animal protecting
*　　　┃　 　 ┗━━━┓ 神兽保佑,代码无bug　
*　　　┃ 　　　　　　 ┣┓
*　　  ┃ 　　　　　 　┏┛
*　    ┗┓┓┏━┳┓┏┛ + + + +
*　　　　┃┫┫　┃┫┫
*　　　　┗┻┛　┗┻┛+ + + +
*/