洛谷传送门

为什么总有人说这是长链剖分板子题？

Solution

题意非常简洁，这让我少了转化题意这一步。

我们考虑什么样的三个点在树上满足两两之间距离 \(d\) 相等：

1. 对于某一个点，它的子树内以它为LCA，距它 \(d\) 的三个点
2. 对于某一个点，它的 \(d\) 级祖先以及子树内两个以它为LCA，距它 \(d\) 的点

那么我们选择用树形DP。

对于情况1，设 \(f_{i,j}\) 表示以 \(i\) 为根的子树中，距 \(i\) 距离为 \(j\) 的点数。

那么情况2怎么表示比较好？看题解发现可以用上面提到的 某一个点 ，即中间点 。

设 \(g_{i,j}\) 表示以 \(i\) 为根的子树中，两个点到LCA的距离相等（此处用 \(d\) 表示），LCA到 \(i\) 的距离为 \(d-j\) 的点对数。

那么可以得到最朴素的转移：

\[ans+=g_{i,0},ans+=\sum_{u,v\in son_i,u\not=v}f_{u,j-1}\times g_{v,j+1}\\g_{i,j}+=\sum_{u,v\in son_i,u\not=v}f_{u,j-1}\times f_{v,j+1}\\ f_{i,j}+=\sum_{u\in son_i}f_{u,j-1},g_{i,j}+=\sum_{u\in son_i}g_{u,j+1} \]

这个转移有的地方比较好想，但是考虑周全还是要谨慎思考的。至于这里面每个转移的含义，画个图会恍然大悟的。

讲完了

因为是最朴素的，在会枚举子树的两个点的转移中，复杂度显然是会爆炸的。

---

来到第一个优化——利用前缀和的思想

考虑将 \(i\) 的一个新儿子 \(u\) 加入树中对 \(ans\) 产生的贡献，应该是 \(ans+=g_{i,j}\times f_{u,j-1}\) ，因为这个儿子只对之前加入的有贡献（也可以自己模拟插入儿子算算）

同理可以得到： \(g_{i,j}+=f_{i,j}\times f_{u,j-1}\)

那么这五种转移就都是 \(O(n)\) 转移的了，时间复杂度就变成了 \(O(n^2)\) (´▽`ʃ♡ƪ)

完结

什么，你说我上面提到了长链剖分？

---

这里是第二次优化——长链剖分

仔细看最后两个转移： \(f_{i,j}+=\sum_{u\in son_i}f_{u,j-1},g_{i,j}+=\sum_{u\in son_i}g_{u,j+1}\)

这样可能不明显，如果只有一个儿子： \(f_{i,j}=f_{son,j-1},g_{i,j}=g_{son,j+1}\)

也就是说，如果只有一个儿子，我们是可以做到直接赋值的

用指针写就是 \(f_i=f_{son}-1,g_{i}=g_{son}+1\)

如果是在树上，我们可以进行长链剖分，让重儿子直接赋值，轻儿子和原来一样转移

因为轻儿子都是重链的顶部，所以转移复杂度就是重链长度

那么全局时间复杂度就是 \(O(\sum len)=O(n)\)

正式完结(≧∇≦)ﾉ

Code

#include<cmath>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#define ll long long

using namespace std;
const int N=100010;
int n;

inline int read(){
    int x=0,f=1;
    char ch=getchar();
    while(!isdigit(ch)){if(ch=='-') f=-1;ch=getchar();}
    while(isdigit(ch)){x=x*10+(ch^48);ch=getchar();}
    return x*f;
}

struct edge{
    int to,nxt;
}e[N<<1];
int head[N],cnt;
inline void add(int u,int v){
    e[++cnt].to=v;
    e[cnt].nxt=head[u];
    head[u]=cnt;
}

int h[N],son[N];
void dfs1(int u,int fa){
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v==fa) continue;
        dfs1(v,u);
        h[u]=max(h[u],h[v]);
        if(h[v]>h[son[u]]) son[u]=v;
    }
    h[u]=h[son[u]]+1;
}

ll *f[N],*g[N],tmp[N<<2],*id=tmp,ans;//这里是指针
void dfs2(int u,int fa){
    if(son[u]) f[son[u]]=f[u]+1,g[son[u]]=g[u]-1,dfs2(son[u],u);
    f[u][0]=1;ans+=g[u][0];
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v==fa||v==son[u]) continue;
        f[v]=id;id+=h[v]<<1;
        g[v]=id;id+=h[v]<<1;
        dfs2(v,u);
        for(int j=0;j<h[v];j++){
            if(j) ans+=f[u][j-1]*g[v][j];
            ans+=g[u][j+1]*f[v][j];
        }
        for(int j=0;j<h[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(){
    n=read();
    for(int i=1,u,v;i<n;i++){
        u=read();v=read();
        add(u,v);add(v,u);
    }
    dfs1(1,0);
    f[1]=id;id+=h[1]<<1;
    g[1]=id;id+=h[1]<<1;
    dfs2(1,0);
    printf("%lld\n",ans);
    return 0;
}