【洛谷】P5904 [POI2014]HOT-Hotels（长链剖分）

原题链接

题意

给出一棵有 $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$ 的距离

对于情况三，对答案的贡献为：

a n s \leftarrow g_{u, 0}

$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}$ 来补偿少的距离。于是这部分对答案的贡献就是：

a n s \leftarrow \sum_{x, y \in s o n_{u}, x \neq y} f_{v, j - 1} \times g_{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} \leftarrow f_{v, j - 1}

$f_{u,j} \gets f_{v,j-1}$

而对于 $g$ 数组的转移，一种是原来就在 $v$ 子树内，即：

g_{u, j} \leftarrow g_{v, j + 1}

$g_{u,j} \gets g_{v,j+1}$

一种是原来有一个点不在 $v$ 子树内，即：

g_{u, j} \leftarrow \sum_{x, y i n s o n_{u}, x \neq y} f_{x, j - 1} \times f_{y, j - 1}

$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;
}