【DP-树形DP】BZOJ3522/BZOJ4543 [POI2014]Hotel

【题目】
简单版
数据加强版
一颗$n$个节点的树。
找三个不同编号的节点，使它们两两间距离相同（一条边距离视作1），求方案数。

【题目分析】
简化版的我们都会做。
由于可以$O(n^2)$搞，很容易想到枚举中间点（树根），然后用树形DP搞一搞
发现三个妹子一定在根的三个不同儿子的子树中，所以我们可以遍历根的每个儿子，设$s[i]$表示其中一个儿子的子树中，深度为i的点有多少个，然后设f[i][j]表示在j个不同儿子的子树中，各选出一个深度为i的点的方案数，可以得到$f[i][j]+=f[i][j-1]*s[i]$
但是到了10w级别，就不能这样了qwq （这不废话吗）

【解题思路】
难度加大以后就变神了啊qwq。%%%题解。
居然是用指针来$O(1)$转移。。。以下来自dalao的blog。

我们考虑一下dp，$f[i][j]$表示以$i$为根的子树里与$i$距离为$j$的点的个数，$g[i][j]$表示子树内有$g[i][j]$对点深度相同，且距离他们的$LCA$距离都为$d$，且$i$与他们的$LCA$的距离为$d-j$。换一种说法是表示以$i$为根的子树里有这么多个点对在底下分叉了，并且还没有第三个点和这个点对匹配，这个第3个点不在$i$的子树里并且与$i$距离为$j$的方案数（不考虑第三个点有多少种选法）

设$x$表示当前点，$y$表示儿子，则$f[x][0]=1,ans+=g[x][0]$

这样的话枚举出边，一边枚举一边更新保证不重复计算，每次枚举出边的时候再枚举$i$

$$f[x][i]+=f[y][i-1]$$

$$g[x][i-1]+=g[y][i]$$
$$g[x][i+1]+=f[x][i+1]*f[y][i]$$
$$ans+=f[x][i-1]*g[y][i]+g[x][i+1]*f[y][i]$$

但是这样的话时间和空间都会爆，我们把整个树进行轻重链剖分，子树深度最大的儿子是重儿子，重边练成重链，对于一个点，在第一次用儿子更新的时候我们有$f[x][i]=f[y][i-1]，g[x][i]=g[y][i+1]$，可以用指针$O(1)$进行这一步转移，由于对一个儿子进行转移的复杂度是$O(dep[y])$，所以不妨对重儿子进行$O(1)$转移

这个时间复杂度是$O(n)$的，证明如下：

设$h[x]$表示以$x$为根的子树的高度
对每个点转移的复杂度为$\sum (h[y]-h[son[x]])=\sum(h[y] -h[x]+1)$，做和的话除了叶子节点所有点的$dep$都被抵消，所以复杂度为$O(n)$
空间的话非叶子节点所需要的空间都是由他所在重链的叶子节点用指针挪过来的，所以对每个叶子节店给他开正比于所在重链长度的空间即可
空间复杂度为$O(\sum 重链长度和)=O(n)$

【参考代码】

BZOJ3522

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

typedef long long LL;
const int N=5010;
int n,tot;
int head[N],dep[N];
LL ans;
LL s[N],f[4][N];

struct Tway
{
    int v,nex;
};
Tway e[N<<1];

void add(int u,int v)
{
    e[++tot]=(Tway){v,head[u]};
    head[u]=tot;
}

void dfs(int x,int fa)
{
    s[dep[x]]++;
    for(int i=head[x];i;i=e[i].nex)
    {
        int v=e[i].v;
        if(v==fa)
            continue;
        dep[v]=dep[x]+1;
        dfs(v,x);
    }
}

int main()
{
//  freopen("BZOJ3522.in","r",stdin);
//  freopen("BZOJ3522.out","w",stdout);

    scanf("%d",&n);
    for(int i=1;i<n;++i)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        add(u,v);add(v,u);
    }
    for(int i=1;i<=n;++i)
    {
        memset(f,0,sizeof(f));
        for(int j=1;j<=n;++j)
            f[0][j]=1;
        for(int j=head[i];j;j=e[j].nex)
        {
            memset(s,0,sizeof(s));
            int v=e[j].v;
            dep[v]=1;dfs(v,i);
            for(int k=3;k;--k)
                for(int q=1;s[q];++q)
                    f[k][q]+=f[k-1][q]*s[q];
        }
        for(int j=1;f[3][j];++j)
            ans+=f[3][j];
    }
    printf("%lld\n",ans);

    return 0;
}

BZOJ4543

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

typedef long long LL;
const int N=1e5+10;

int n,tot;
int head[N],son[N],dep[N],h[N],mx[N],siz[N];
LL id,ans;
LL mem[10*N],*f[N],*g[N];

struct Tway
{
    int v,nex;
};
Tway e[N<<1];

void add(int u,int v)
{
    e[++tot]=(Tway){v,head[u]};
    head[u]=tot;
}

void dfs(int u,int fa)
{
    dep[u]=dep[fa]+1;mx[u]=u;
    for(int i=head[u];i;i=e[i].nex)
    {
        int v=e[i].v;
        if(v==fa)
            continue;
        dfs(v,u);
        if(dep[mx[v]]>dep[mx[son[u]]])
            son[u]=v;
        if(dep[mx[v]]>dep[mx[u]])
            mx[u]=mx[v];
    }
    siz[mx[u]]=h[u]=dep[mx[u]]-dep[u]+1;
}

void init()
{
    for(int i=1;i<=n;++i)
    {
        if(mx[i]!=i)
            continue;
        id+=siz[i]-1;
        f[i]=&mem[id];++id;
        g[i]=&mem[id];id+=2*siz[i]+2;
    }
}

void solve(int u,int fa)
{
    if(mx[u]==u)
    {
        f[u][0]=1;
        return;
    }
    solve(son[u],u);
    f[u]=f[son[u]]-1;g[u]=g[son[u]]+1;
    ans+=g[u][0];f[u][0]=1;

    for(int i=head[u];i;i=e[i].nex)
    {
        int v=e[i].v;
        if(v==fa || v==son[u])
            continue;
        solve(v,u);

        for(int i=0;i<=h[v];++i)
            ans+=f[v][i]*g[u][i+1];
        for(int i=1;i<=h[v];++i)
            ans+=g[v][i]*f[u][i-1];

        for(int i=0;i<=h[v];++i)
            g[u][i+1]+=f[v][i]*f[u][i+1];
        for(int i=1;i<=h[v];++i)
            g[u][i-1]+=g[v][i];

        for(int i=0;i<=h[v];++i)
            f[u][i+1]+=f[v][i];
    }
}

int main()
{
    freopen("BZOJ4543.in","r",stdin);
    freopen("BZOJ4543.out","w",stdout);

    scanf("%d",&n);
    for(int i=1;i<n;++i)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        add(u,v);add(v,u);
    }
    dfs(1,0);
    init();
    solve(1,0);
    printf("%lld\n",ans);

    return 0;
}

【总结】
至今仍觉得这个时间复杂度的证明很妙啊！
然后这个指针的$O(1)$转移也是真神啊！