CF337D——换根dp

传送门

题意

一棵树，有些是红点。如果点u满足：每个红点到u的距离都<=D，则u符合题意。求合题意的点有多少个。

注：题干用的是”小姐姐“，本文都用”红点“来指代。

思路

我们不妨求出离当前点最远的红点的距离值，这个距离<=D即符合题意。这是一个全局性的统计问题，所以我们很快能想到这是一个换根dp。定义down[u]是以u为子树的最远距离，up[u]是u子树以外的点的最远距离。up既可以定义为包含u，也可以定义为不包含。之所以up不定义为全局的，是因为max运算符不能做差。

答案是max(down[u],up[u])。

down的转移：如果v子树存在至少一个红点，则更新down[u] = max(down[u],down[v]+1)，否则不更新。如果所有v子树都没有红点，则down[u]=0。

up的转移：如果u子树以外的点（u不在统计范围内）没有红点，则up[u]=0。否则我们看着图来求：

我们得到：up[u] = 1+max(up[ufa],ufa除u以外的孩子的红点到ufa的最大距离)。

因为max不能做差，所以需要维护一个前缀max和后缀max来合成。于是我们定义了lef[u][idx]是0~idx号孩子的down[v]的最大值，rig同理。

所以ufa除u以外的孩子的红点到ufa的最大距离是：max(!idx ? 0 : 1+lef[ufa][idx-1],idx+1 >= G[ufa].size() ? 0 : 1+rig[ufa][idx+1])。

注意：记得特判u为根的情况！

最后，我们需要判定”u子树以外的点没有红点“，所以需要定义sz_down表示u子树的红点个数，sz_up表示u子树以外的点的红点个数（这里我们都定义为包含u的，定义为不包含也行）。sz_up[u] = (sz_up[ufa] - a[ufa]) + (sz_down[ufa] - sz_down[u]) + a[u]。

一发AC真爽！

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
#define rep(i,a,b) for(int i = (a);i <= (b);++i)
#define re_(i,a,b) for(int i = (a);i < (b);++i)
#define dwn(i,a,b) for(int i = (a);i >= (b);--i)

const int N = 1e5 + 5;

int n,m,D;bool a[N];
vector<int> G[N],lef[N],rig[N];
int sz_down[N],sz_up[N],down[N],up[N];

void dbg(){puts("");}
template<typename T, typename... R>void dbg(const T &f, const R &... r) {
    cout << f << " ";
    dbg(r...);
}
template<typename Type>inline void read(Type &xx){
    Type f = 1;char ch;xx = 0;
    for(ch = getchar();ch < '0' || ch > '9';ch = getchar()) if(ch == '-') f = -1;
    for(;ch >= '0' && ch <= '9';ch = getchar()) xx = xx * 10 + ch - '0';
    xx *= f;
}

void dfs1(int u,int ufa){
    sz_down[u] = a[u];
    for(int v: G[u]){
        if(v == ufa) continue;
        dfs1(v,u);
        sz_down[u] += sz_down[v];
        if(sz_down[v]) down[u] = max(down[u],down[v]+1);
    }
    const int chNum = G[u].size();
    lef[u] = vector<int>(chNum,0);
    rig[u] = vector<int>(chNum,0);
    re_(i,0,chNum){
        lef[u][i] = max(i ? lef[u][i-1] : 0,down[G[u][i]]);
    }
    dwn(i,chNum-1,0){
        rig[u][i] = max(i < chNum-1 ? rig[u][i+1] : 0,down[G[u][i]]);
    }
}

void dfs2(int u,int ufa,int idx = 0){
    if(ufa){
        sz_up[u] = (sz_up[ufa] - a[ufa]) + (sz_down[ufa] - sz_down[u]) + a[u];
        up[u] = !(sz_up[u] - a[u]) ? 0 : (1 + max(up[ufa],
            max(!idx ? 0 : 1+lef[ufa][idx-1],idx+1 >= G[ufa].size() ? 0 : 1+rig[ufa][idx+1])
        ));
    }
    else sz_up[u] = a[u];
    re_(i,0,G[u].size()){
        int v = G[u][i];
        if(v == ufa) continue;
        dfs2(v,u,i);
    }
}

void dbg1(){
    rep(i,1,n) cout << down[i] << " \n"[i == n];
    rep(i,1,n) cout << sz_down[i] << " \n"[i == n];
    rep(i,1,n){
        dbg(i,":");
        for(int v: lef[i]) cout << v << " ";puts("");
        for(int v: rig[i]) cout << v << " ";puts("");
    }
    puts("");
}

void dbg2(){
    rep(i,1,n) cout << up[i] << " \n"[i == n];
    rep(i,1,n) cout << sz_up[i] << " \n"[i == n];
    rep(i,1,n) if(max(down[i],up[i]) <= D) cout << i << " ";puts("");
}

int main(int argc, char** argv) {
    read(n);read(m);read(D);
    rep(_,1,m){
        int x;read(x);a[x] = true;
    }
    re_(i,1,n){
        int x,y;read(x);read(y);
        G[x].push_back(y);G[y].push_back(x);
    }
    dfs1(1,0);
    // dbg1();
    dfs2(1,0);
    // dbg2();
    int ans = 0;
    rep(i,1,n) ans += (max(down[i],up[i]) <= D);
    printf("%d\n",ans);
    return 0;
}