题意
给出一棵有 \(n\) 个点的树,求有多少组点 \((i,j,k)\) 满足 \(i,j,k\) 两两之间的距离都相等。
\((i,j,k)\) 与 \((i,k,j)\) 算作同一组。
\(1\le n\le10^5\)。
思路
对于每一个节点 \(u\),统计其子树内的答案时,仅 \((i,j,k)\) 满足 \(\mathrm{LCA}(i,j,k)=u\) 的情况,这样就可以避免重复计算。而对于一棵子树内选择三个节点的情况,又可以分为若干种。
第一种情况: \(\mathrm{LCA}(i,j)=\mathrm{LCA}(j,k)=\mathrm{LCA}(i,k)=u\)。
第二种情况: \(\mathrm{LCA}(i,j)=\mathrm{LCA}(i,k)=u \ne \mathrm{LCA}(j,k)\)。
第三种种情况: \(\mathrm{LCA}(j,k) \ne i=u\)。
不难发现,任意一种情况都要求两个点到其 \(\mathrm{LCA}\) 的距离相等。
于是可以设 \(f_{u,k}\) 为以 \(u\) 为根的子树中,到 \(u\) 距离为 \(k\) 的点的数量。
而由于还有第二、三种情况的存在,我们需要维护一个 \(g_{u,x}\),表示在子树 \(u\) 中,点对 \((x,y)\) 到 \(LCA\) 的距离比 \(u\) 到 \(LCA\) 的距离大 \(x\) 的总数。即 \(dis(u,LCA)=dis(v,LCA)=dis(u,LCA)+x\)。其中当 \(LCA=u\) 时,\(x\) 就表示 \(x\) 或 \(y\) 到 \(u\) 的距离
对于情况三,对答案的贡献为:
\[ans \gets g_{u,0} \]对于情况二,可以枚举一棵子树 \(g_{v,j+1}\),那么也就是 \(dis(x,LCA)\) 比 \(dis(v,LCA)\) 大了 \(j+1\),也就是比 \(dis(u,LCA)\) 大了 \(j\),于是就需要枚举另一个子树中的 \(f_{v,j-1}\) 来补偿少的距离。于是这部分对答案的贡献就是:
\[ans \gets \sum_{x,y\in son_u,x \ne y} f_{v,j-1} \times g_{v,j+1} \]注意到 \(g_{u,j}\) 其实也包括了 \(LCA(x,y)=u,dis(x,u)=j\) 的情况,所以这个转移也把情况一包含进去了。
对于 \(f\) 的转移,比较显然:
\[f_{u,j} \gets f_{v,j-1} \]而对于 \(g\) 数组的转移,一种是原来就在 \(v\) 子树内,即:
\[g_{u,j} \gets g_{v,j+1} \]一种是原来有一个点不在 \(v\) 子树内,即:
\[g_{u,j} \gets \sum_{x,y in son_u,x \ne y} f_{x,j-1} \times f_{y,j-1} \]注意到转移都是 \(j\) 从 \(0\) 开始的,所以可以用前缀和优化。而 \(j\) 实际上和树的深度有关系,于是可以考虑长链剖分优化。\(f_u,g_u\) 和 \(f_v,g_v\) 的转移有很大一部分只需要更改一下下标,于是可以用到动态分配内存的小技巧。总的时间复杂度就为 \(O(n)\)。
code:
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<iostream>
using namespace std;
const int N=1e5+10;
#define LL long long
LL *f[N],*g[N],tmp[N<<2],*p=tmp,ans;
int n,h[N],son[N],dep[N],idx;
struct edge{int v,nex;}e[N<<1];
void add(int u,int v){e[++idx]=edge{v,h[u]};h[u]=idx;}
void dfs(int u,int fa)
{
for(int i=h[u];i;i=e[i].nex)
{
int v=e[i].v;if(v==fa) continue;dfs(v,u);
if(!son[u]||dep[v]>dep[son[u]]) son[u]=v;
}
dep[u]=dep[son[u]]+1;
}
void solve(int u,int fa)
{
if(son[u]) f[son[u]]=f[u]+1,g[son[u]]=g[u]-1,solve(son[u],u);
f[u][0]=1,ans+=g[u][0];
for(int i=h[u],v;i;i=e[i].nex) if((v=e[i].v)!=fa&&v!=son[u])
{
f[v]=p,p+=dep[v]*2;g[v]=p,p+=dep[v]*2;solve(v,u);
for(int i=0;i<dep[v];i++)
{
if(i) ans+=f[u][i-1]*g[v][i];
ans+=g[u][i+1]*f[v][i];
}
for(int i=0;i<dep[v];i++)
{
g[u][i+1]+=f[u][i+1]*f[v][i];
if(i) g[u][i-1]+=g[v][i];
f[u][i+1]+=f[v][i];
}
}
}
int main()
{
scanf("%d",&n);for(int u,v,i=1;i<n;i++) scanf("%d%d",&u,&v),add(u,v),add(v,u);dfs(1,0);
f[1]=p;p+=dep[1]*2;g[1]=p;p+=dep[1]*2;solve(1,0);printf("%lld\n",ans);
return 0;
}
标签:长链,gets,洛谷,剖分,int,son,LCA,dis,mathrm
From: https://www.cnblogs.com/ListenSnow/p/17230972.html