bzoj 4543 HOTEL 加强版

题目大意：

求树上取三个点这三个点两两距离相等的方案数

思路：

远古时候的$n^2$做法是换根 但那样无法继续优化了

学习了一波长链剖分

考虑如何在一棵树上进行dp

设$f[i][j]$表示以$i$为根的子树内与$i$的距离为$j$的点数量

$g[i][j]$表示以$i$为根的子树内满足与lca距离为$d$且lca与$i$的距离为$d-j$的点对数（lca在子树内）

对于每个子树 对答案的贡献为$g[x][0]$与对每一个新进来的子树$\sum _{i=0} ^{mxd[v]} f[x][i-1]*g[v][i]+g[x][i+1]*f[v][i]$

更新$f,g$的时候除了正常的继承的子树 $g[x][i+1]+=f[x][i+1]*f[v][i]$

 #include<iostream>
 #include<cstdio>
 #include<cstring>
 #include<cstdlib>
 #include<cmath>
 #include<algorithm>
 #include<queue>
 #include<vector>
 #include<map>
 #include<set>
 #define ll long long
 #define inf 2139062143
 #define MAXN 100100
 #define MOD 998244353
 #define rep(i,s,t) for(register int i=(s),i##__end=(t);i<=i##__end;++i)
 #define dwn(i,s,t) for(register int i=(s),i##__end=(t);i>=i##__end;--i)
 #define ren for(register int i=fst[x];i;i=nxt[i])
 #define pb(i,x) vec[i].push_back(x)
 #define pls(a,b) (a+b)%MOD
 #define mns(a,b) (a-b+MOD)%MOD
 #define mul(a,b) (1LL*(a)*(b))%MOD
 using namespace std;
 inline int read()
 {
     int x=,f=;char ch=getchar();
     while(!isdigit(ch)) {if(ch=='-') f=-;ch=getchar();}
     while(isdigit(ch)) {x=x*+ch-'';ch=getchar();}
     return x*f;
 }
 int n,*f[MAXN],*g[MAXN],tmp[MAXN<<],*id=tmp,ans;
 int fst[MAXN],nxt[MAXN<<],to[MAXN<<],cnt,mxd[MAXN],son[MAXN];
 void add(int u,int v) {nxt[++cnt]=fst[u],fst[u]=cnt,to[cnt]=v;}
 void dfs(int x,int pa)
 {
     ren if(to[i]^pa) {dfs(to[i],x);if(mxd[to[i]]>mxd[son[x]]) son[x]=to[i];}
     mxd[x]=mxd[son[x]]+;
 }
 void New(int x) {f[x]=id,id+=mxd[x]<<,g[x]=id,id+=mxd[x]<<;}
 void dp(int x,int pa)
 {
     if(son[x]) {f[son[x]]=f[x]+,g[son[x]]=g[x]-;dp(son[x],x);}
     f[x][]=,ans+=g[x][];
     ren if(to[i]^pa&&to[i]^son[x])
     {
         New(to[i]);dp(to[i],x);
         rep(j,,mxd[to[i]])
         {
             ans+=g[x][j+]*f[to[i]][j];
             if(j) ans+=f[x][j-]*g[to[i]][j];
         }
         rep(j,,mxd[to[i]])
         {
             g[x][j+]+=f[x][j+]*f[to[i]][j],f[x][j+]+=f[to[i]][j];
             if(j) g[x][j-]+=g[to[i]][j];
         }
     }
 }
 int main()
 {
     n=read();int a,b;rep(i,,n) a=read(),b=read(),add(a,b),add(b,a);
     dfs(,);New();dp(,);printf("%d\n",ans);
 }