【BZOJ4543】【POI2014】Hotel加强版（长链剖分）

传送门

题意：求树上满足三点之间距离两两相等的三元组个数

$n\le 1e5$

原题数据是$n\le 5000$

考虑怎么做
$f[u][i]$表示$u$为根，深度为$i$的点的个数
$g[u][i]$表示$u$为根，满足2点到$lca$的距离减去$lca$到$u$的距离为$i$，即dep[x]+dep[y]−3∗deplca=idep[x]+dep[y]-3*dep_{lca}=idep[x]+dep[y]−3∗deplca​=i的点对个数
换句话说就是还差$i$个距离满足能凑成$3$元组的点对个数

则
$ans+=g[u][i+1]\ast f[v][i];$
$ans+=f[u][i-1]\ast g[v][i];$
$g[u][i+1]+=f[u][i+1]\ast f[v][i];$
$f[u][i+1]+=f[v][i];$
$g[u][i-1]+=g[v][i];$

这式子很显然吧
发现转移的时候
$f[u][i+1]+=f[v][i];$
$g[u][i-1]+=g[v][i];$
既然只和深度有关，
就可以愉快的长链剖分了

复杂度$O(n)$
据说可以点分$O(nlogn)$
关我p事

#include<bits/stdc++.h>
using namespace std;
const int RLEN=1<22|1;
#define ll long long
inline char gc(){static char ibuf[RLEN],*ob,*ib;(ob==ib)&&(ob=(ib=ibuf)+fread(ibuf,1,RLEN,stdin));return (ib==ob)?EOF:*ib++;
}
inline int read(){char ch=gc();int res=0,f=1;while(!isdigit(ch)){if(ch=='-')f=-f;ch=gc();}while(isdigit(ch))res=(res+(res<<2)<<1)+(ch^48),ch=gc();return res*f;
}
const int N=1000005;
ll *f[N],*g[N],*id,tmp[N<<2],ans;
int n,adj[N],nxt[N<<1],to[N<<1],dep[N],son[N],cnt;
inline void addedge(int u,int v){nxt[++cnt]=adj[u],adj[u]=cnt,to[cnt]=v;
}
void dfs1(int u,int fa){for(int e=adj[u];e;e=nxt[e]){int v=to[e];if(v==fa)continue;dfs1(v,u);if(dep[v]>dep[son[u]])son[u]=v;}dep[u]=dep[son[u]]+1;
}
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 e=adj[u];e;e=nxt[e]){int v=to[e];if(v==fa||v==son[u])continue;f[v]=id,id+=dep[v],g[v]=id+dep[v],id+=dep[v]*2;dfs2(v,u);for(int i=dep[v]-1;~i;i--){ans+=g[u][i+1]*f[v][i];if(i)ans+=f[u][i-1]*g[v][i];g[u][i+1]+=f[u][i+1]*f[v][i];f[u][i+1]+=f[v][i];}for(int i=dep[v]-1;i;i--){g[u][i-1]+=g[v][i];}}
}
int main(){n=read();for(int i=1;i<n;i++){int u=read(),v=read();addedge(u,v),addedge(v,u);}dfs1(1,0);id=tmp;f[1]=id,id+=dep[1],g[1]=id+dep[1],id+=dep[1]*2;dfs2(1,0);cout<<ans;
}