正题

题目链接:https://www.luogu.com.cn/problem/P5904

---

题目大意

$n$个点的一棵树，求有多少个点对$(i,j,k)$使得这三个点距离相等。

---

解题思路

有两种情况，一是$i$是$j,k$的祖先，二是$i,j,k$互相没有祖先关系

考虑$dp$，$f_{i,j}$表示$i$点的子树中与$i$距离为$j$的点的个数，然后$g_{i,j}$表示$i$的子树中满足有多少个点$(x,y)$对使得$dis(x,lca)=dis(y,lca)=dis(i,lca)+j$。然后考虑统计答案
$$ans+=g_{y,i}\ast f_{x,i-1}+g_{x,i+1}\ast f_{y,i}$$
然后考虑这个点对$g$的影响有
$$g_{x,i+1}+=f_{x,i+1}\ast f_{y,i}$$
然后继承$g_{x,i-1}+=g_{y,i},f_{x,i+1}+=f_{y,i}$

然后可以长链剖分进行优化，$f$在长链上向后运动，但是$g$是向前运动的。

考虑结点$1$，因为每次向前运动要放在$le{n}_1$的位置，但是该位置向后又要储存$le{n}_1$个值，所以我们要开两倍的空间

时间复杂度$O(n)$

---

$code$

#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
const ll N=4e5+10;
struct node{ll to,next;
}a[N*2];
ll n,tot,ans,ls[N],son[N];
ll len[N],buff[N*2],bufg[N*2];
ll *f[N*2],*g[N*2],*nowf,*nowg;
void addl(ll x,ll y){a[++tot].to=y;a[tot].next=ls[x];ls[x]=tot;return;
}
void dfs(ll x,ll fa){for(ll i=ls[x];i;i=a[i].next){ll y=a[i].to;if(y==fa)continue;dfs(y,x);if(len[y]>len[son[x]])son[x]=y;}len[x]=len[son[x]]+1;return;
}
void solve(ll x,ll fa){f[x][0]=1;if(son[x]){f[son[x]]=f[x]+1;g[son[x]]=g[x]-1;solve(son[x],x);}ans+=g[x][0];for(ll i=ls[x];i;i=a[i].next){ll y=a[i].to;if(y==son[x]||y==fa)continue;f[y]=nowf;nowf+=len[y];nowg+=len[y]*2+10;g[y]=nowg++;solve(y,x);for(ll j=0;j<len[y];j++){if(j)ans+=g[y][j]*f[x][j-1];ans+=g[x][j+1]*f[y][j];}for(ll j=0;j<len[y];j++){g[x][j+1]+=f[y][j]*f[x][j+1];if(j)g[x][j-1]+=g[y][j];f[x][j+1]+=f[y][j];}}return;
}
int main()
{scanf("%lld",&n);for(ll i=1;i<n;i++){ll x,y;scanf("%lld%lld",&x,&y);addl(x,y);addl(y,x);}dfs(1,1);nowf=buff;nowg=bufg;nowf+=len[1];nowg+=len[1]*2+10;f[1]=buff;g[1]=nowg++;solve(1,1);printf("%lld",ans);
}