4543: [POI2014]Hotel加强版[树形DP+长链剖分]

题目链接:https://www.lydsy.com/JudgeOnline/problem.php?id=4543

解题思路:

长链剖分

定义:

f[i][j]表示以i为根节点的子树，有多少个节点和i的距离是j的.

g[i][j]表示以i为根节点的子树，在子树外一个距离i为j的点可以跟i子树内的两个点组成两两相等的方案数.

那么就有:

f[u][j+1] += f[v][j],  g[u][j+1] += f[u][j+1]*f[v][j],  g[u][j] += g[u][j+1].

所以当u节点继承重儿子son[u]的时候有:f[u][j+1] = f[son[u]][j],g[u][j] = g[son[u]][j+1].

如果我们用指针维护他们的数组首地址,假设f[u] = a,那么f[son[u]] = a+1,假设g[u] = b,那么g[son[u]] = b - 1;

因为长链剖分的性质，所有长链的点长度和==n，所以可以用指针在一维数组上面维护二维数组。

那么ans就可以用f[u][j]*g[v][j+1]和g[u][j+1]*f[v][j]更新，需要注意的是u节点从重儿子继承过来还需加上g[u][0]

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int mx = 1e5 + 10;
vector <int> r[mx];
int n,dep[mx],son[mx];
ll tmp[mx<<2],*id = tmp,*f[mx],*g[mx];
ll ans;
void dfs(int u,int fa){
    for(int i=0;i<r[u].size();i++){
        int v = r[u][i];
        if(v==fa) continue;
        dfs(v,u);
        if(dep[v]>dep[son[u]]) son[u] = v;
    }
    dep[u] = dep[son[u]] + 1; 
}
void dp(int u,int fa){
    if(son[u]){
        f[son[u]] = f[u] + 1;
        g[son[u]] = g[u] - 1;
        dp(son[u],u);
    }
    f[u][0] = 1;
    ans += g[u][0];//以u做最高点的方案数，这里处理重儿子的时候可能被忽略
    for(int i=0;i<r[u].size();i++){
        int v = r[u][i];
        if(v==fa||v==son[u]) continue;
        f[v] = id,id += dep[v]<<1;
        g[v] = id,id += dep[v]<<1;
        dp(v,u);
        for(int j=0;j<dep[v];j++){
            ans += g[u][j+1]*f[v][j];
            if(j) ans += f[u][j-1]*g[v][j];
        }
        for(int j=0;j<dep[v];j++){
            g[u][j+1] += f[u][j+1]*f[v][j];
            if(j) g[u][j-1] += g[v][j];
            f[u][j+1] += f[v][j];
        }
    } 
}
int main(){
    scanf("%d",&n);
    int u,v;
    for(int i=1;i<n;i++){
        scanf("%d%d",&u,&v);
        r[u].push_back(v);
        r[v].push_back(u);
    }
    dfs(1,0);
    f[1] = id,id += dep[1]<<1;
    g[1] = id,id += dep[1]<<1;
    dp(1,0);
    printf("%lld\n",ans);
    return 0;
}