Problem
有一颗有 n n n 个点的树,边权为 1 1 1。
求满足 d i s x , y = d i s y , z = d i s x , z dis_{x,y}=dis_{y,z}=dis_{x,z} disx,y=disy,z=disx,z 的无序三元组 ( x , y , z ) (x,y,z) (x,y,z) 的数量( x , y , z x,y,z x,y,z 两两不同)。
数据范围: n ≤ 1 0 5 n≤10^5 n≤105。
Solution
先贴一篇题解吧,明天再写。
https://blog.csdn.net/qq_42555009/article/details/88917806
Code
#include<cstdio>
#include<cstring>
#include<algorithm>
#define N 100005
#define ll long long
using namespace std;
int n,t;
int first[N],v[N<<1],nxt[N<<1];
int son[N],len[N];
ll ans,temp[N<<2],*f[N],*g[N],*pos=temp;
void add(int x,int y){
nxt[++t]=first[x],first[x]=t,v[t]=y;
}
void dfs(int x,int fa){
for(int i=first[x];i;i=nxt[i]){
int to=v[i];
if(to==fa) continue;
dfs(to,x);
if(len[to]>len[son[x]]) son[x]=to;
}
len[x]=len[son[x]]+1;
}
void dp(int x,int fa){
if(son[x]) f[son[x]]=f[x]+1,g[son[x]]=g[x]-1,dp(son[x],x);
f[x][0]=1,ans+=g[x][0];
for(int i=first[x];i;i=nxt[i]){
int to=v[i];
if(to==fa||to==son[x]) continue;
f[to]=pos,pos+=len[to],g[to]=pos+len[to],pos+=2*len[to];
dp(to,x);
for(int j=len[to]-1;~j;--j){
ans+=g[x][j+1]*f[to][j];
if(j) ans+=f[x][j-1]*g[to][j];
g[x][j+1]+=f[x][j+1]*f[to][j];
f[x][j+1]+=f[to][j];
}
for(int j=len[to]-1;j;--j) g[x][j-1]+=g[to][j];
}
}
int main(){
int x,y,i;
scanf("%d",&n);
for(i=1;i<n;++i){
scanf("%d%d",&x,&y);
add(x,y),add(y,x);
}
dfs(1,0);
f[1]=pos,pos+=len[1],g[1]=pos+len[1],pos+=2*len[1];
dp(1,0);
printf("%lld",ans);
return 0;
}