洛谷P5904 树形dp，长链剖分优化

题意：

给出一棵树，在树上选三个点，使得他们两两树上距离相等，问有多少种选法？

方法：

能成立的点对应该是如图这样的关系

选出的三个点是$(x,y,z)$，他们由一个中转点$T$中转，使得$XT=YT=ZT=k$，即三个点分别从$T$的三棵子树中选取

但由于树的特性，至少有一支成为父亲，暴力枚举子树的时间复杂度是难以接受的

于是考虑树形$dp$，设$f[u][i]$代表在$u$的子树内距离$u$为$i$的点有多少个，$g[u][i]$代表在$u$的子树内，存在多少对二元组$(x,y)$，令$lca$为$x,y$的最近公共祖先，有$dis(x,lca)=dis(y,lca)$并且$dis(u,lca)=dis(x,lca)-i$成立。

为什么这样考虑？可以发现这样设置$g$后，$x,y$被确定了，那么只需要去寻找$z$，而$lca$距离$u$是$dis(x,lca)-i$，我们要找到另外子树的一个点$z$，使$dis(z,lca)=dis(x,lca)$，那么只需要从$u$出发，找到另外子树的距离$u$长度为$i$的点，即$f[u][i]$。

不妨将$g[u][i]$解释为一个等价形式，意为在$u$的子树内，存在多少点对$(x,y)$，使得其他子树内距离$u$为$i$的点都可以成为$z$

并且由于他们以深度为下标，那么可以考虑长链剖分来合并$dp$

转移方法：

$$ans+=g[u][0]+\sum _{v\in su{m}_u}g[v][i]\ast f[u][i-1]+g[u][i+1]\ast f[v][i]$$

$$g[u][i-1]+=\sum _{v\in su{m}_u}g[v][i]\text{\hspace{1em}(1)}$$

$$g[u][i+1]+=f[v][i]\ast f[u][i+1]\text{\hspace{1em}(2)}$$

$$f[u][i+1]=f[v][i]$$

以上的$i$都是枚举$v$的长链长，因此$j\in [0,max1[v]-1]$

$ans:$

答案的计算不妨考虑结果点$z$在哪里？如果恰好在$u$，那么答案即$g[u][0]$，如果不在$u$，在另外的子树上，那么$ans+={\sum }_{i=0}^{max1[v]-1}g[u][i]\ast f[u][i]$，如果在当前子树上，他也会被$(x,y)$在的子树比较小的时候归入不在那个时候的子树的情况，不需要重复计数。

$g:$

$g$的转移$(1)$可以理解为，到达了$v$还需要$i$的长度才能到达$z$，那么在$u$还需要的长度自然$-1$，于是$g[u][i-1]=g[v][i]$

转移$(2)$，由于转移$u$的时候，所有以$u$的儿子为$lca$的情况都被添加到$g[u]$了，那么只需要考虑以$u$为$lca$的情况，此时如果$u$为$lca$，那么按照定义，$dis(x,lca)-i=dis(lca,u)=0$，于是$dis(x,lca)=dis(y,lca)=i$，那么$g[u][i]$就是添加有多少点对$(x,y)$，满足不在$u$的同一颗子树内，距离$u$距离都为$i$的点，这样按照前缀和合并子树的思想就可以计算出来

转移顺序：

显然我们统计答案是用当前子树的答案和之前所有子树的答案计算，于是计算$ans$需要放在转移$f,g$之前，这是合并子树前缀和计算的思想。

转移式中计算$ans$是还需要加上$g[u][0]$，这项必须在计算完长儿子后，统计短儿子前加入，如果在最后加入，那么$g[u][0]$包含了$g[v][1]$，$g[v][1]$又包含了$g[vv][2]$，而$g[v][1]$已经和$f[u][1]$计数了，所以不能在最后计数

指针分配内存：

按照转移式，$f$和$g(1)$是可以$O(1)$继承长儿子的，对$g$，$g[son[u]][i]$被写在$g[u][i-1]$上，那么$g[son[u]]=g[u]-1$

，同理$f[son[u]]=f[u]+1$，两个指针偏移方向不一样。一种解决方法是，只用一个数组供$f,g$一起填写，每次分配多一倍的空间，保证有足够空间可以写。先分配$f$再分配$g$的顺序不能调换，除非被写数组的开始分配的点在末尾

        f[v]=pp; pp+=max1[v]<<1;
        g[v]=pp; pp+=max1[v]<<1;

总代码：

#include
#define ll long long
using namespace std;

struct way
{
    int to,next;
}edge[200005];
int cntt,head[100005];

void add(int u,int v)
{
    edge[++cntt].to=v;
    edge[cntt].next=head[u];
    head[u]=cntt;
}

int n,max1[100005],son[100005],depth[100005];

void dfs1(int u,int fa)
{
    for(int i=head[u];i;i=edge[i].next)
    {
        int v=edge[i].to;
        if(v==fa) continue;
        dfs1(v,u);
        if(max1[v]>max1[son[u]]) son[u]=v;
    }
    max1[u]=max1[son[u]]+1;
}

ll ans=0,*f[200005],*g[200005],tmp1[400005],*pp=tmp1;

void dfs(int u,int fa)
{
    if(son[u])
    {
        f[son[u]]=f[u]+1;
        g[son[u]]=g[u]-1;
        dfs(son[u],u);
    }
    f[u][0]=1; ans+=g[u][0];
    for(int i=head[u];i;i=edge[i].next)
    {
        int v=edge[i].to;
        if(v==fa||v==son[u]) continue;
        f[v]=pp; pp+=max1[v]<<1;
        g[v]=pp; pp+=max1[v]<<1;
        dfs(v,u);
        for(int j=0;j<max1[v];j++)
        {
            ans+=f[v][j]*g[u][j+1];
            if(j) ans+=g[v][j]*f[u][j-1];
        }
        for(int j=0;j<max1[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()
{
    cin>>n;
    for(int i=1;i<n;i++)
    {
        int u,v; scanf("%d%d",&u,&v);
        add(u,v); add(v,u);
    }
    dfs1(1,0);
    f[1]=pp; pp+=max1[1]<<1;
    g[1]=pp; pp+=max1[1]<<1;
    dfs(1,0);
    cout<<ans;
    return 0;
;
    g[1]=pp; pp+=max1[1]<<1;
    dfs(1,0);
    cout<<ans;
    return 0;
}