OI - HOT Hotels 加强版(洛谷)

经常口胡长链剖分 $O(n)$，现在算是见识了

---

# 题面

> Linked 洛谷 P5904

点击展开/折叠 题面

给出一棵有 $n$ （$1\le n\le10^5$）个点的树，求有多少组无序点对 $(i,j,k)$ 满足 $i,j,k$ 两两之间的距离都相等。

---

# 解析

三个点 $(x,y,z)$ 两两之间的距离一定可以表示成下面这种形式：

$$\begin{cases} \text{dist}(x,y)=\text{dist}(x,\text{lca}(x,y))+\text{dist}(y,\text{lca}(x,y))\\ \text{dist}(y,z)=\text{dist}(y,\text{lca}(x,y))+\text{dist}(\text{lca}(x,y),z)\\ \text{dist}(z,x)=\text{dist}(x,\text{lca}(x,y))+\text{dist}(\text{lca}(x,y),z) \end{cases}$$

即下面这种情况：

然后因为两两距离相等，可以推得：

令 $\text{dist}(x,lca)=\text{dist}(y,lca)=i$ 则有

$$\text{dist}(x,y)=2i\\\text{dist}(y,z)=\text{dist}(z,x)=i+\text{dist}(lca,z)$$

于是 $\text{dist}(lca,z)=i$。

所以我们可以DP维护两个值：

• $f_{u,i}$：$u$ 子树内与 $u$ 距离为 $i$ 的点数；
• $g_{u,i}$：$u$ 子树内的点对 $(v,w)$ 的数量，满足：
  • $\text{dist}(lca,v)=\text{dist}(lca,w)$；
  • $\text{dist}(lca,u)=\text{dist}(lca,v)-i$。（这样设计方便计算——可以直接把 $g_{u,i}$ 和 $f_{u,i}$ 配对得到答案）

于是树形DP过程中，可以在合并子树时计算答案。设 $u$ 要合并子树 $v$，则

$$ans'=ans+\sum f_{u,i}g_{v,i}+f_{v,i}g_{u,i}\\ f'_{u,i}=f _{u,i}+f _{v,i-1}\\ g'_{u,i}=g _{u,i}+f _{u,i}f _{v,i-1}$$

接下来考虑优化，发现 $f,g$ 第二维的取值范围均是 $[0,dep_u)$（这里的深度定义是该子树最深的节点的相对深度），于是可以长链剖分优化。

根据长链剖分的原理，每个点恰好属于一条长链，且每条长链只会在链顶作为“轻儿子”暴力合并的时候会有链长的复杂度，链长和是 $n$，则复杂度 $O(n)$。

然后就是一些小技巧，比如分配内存池，用指针比较快乐，只是注意分配的 siz（设 $v$ 是 $u$ 的重儿子）：

1. 对于 $f_u$，$f_{u,i+1}:=f_{v,i},f_{u,0}=1$，则可以直接把 $f_v$ 的初始指针指向 $f_u+1$，一条长链只需要链长的空间；
2. 对于 $g_u$，$g_{u,i}:=g_{v,i+1}$，因此需要把 $g_{v}$ 的初始指针指向 $g_{u}-1$，这样所需空间是链长两倍的。

---

# 源代码

/*Lucky_Glass*/
#include<cstdio>
#include<cstring>
#include<cassert>
#include<algorithm>
using namespace std;

inline int Read(){
	register int r=0,c=getchar();
	while(c<'0' || '9'<c) c=getchar();
	while('0'<=c && c<='9') r=(r<<1)+(r<<3)+(c^'0'),c=getchar();
	return r;
}

typedef long long ll;
#define ci const int &
const int N=1e5+10;
#define alloc(x) (f[x]=fcnt,fcnt+=dep[x],g[x]=(gcnt+=dep[x]),gcnt+=dep[x])

struct GRAPH{
	int head[N],to[N<<1],nxt[N<<1],ncnt;
	void AddEdge(ci u,ci v){
		int p=++ncnt,q=++ncnt;
		to[p]=v,nxt[p]=head[u],head[u]=p;
		to[q]=u,nxt[q]=head[v],head[v]=q;
	}
	inline int operator [](ci u){return head[u];}
}Gr;

int n;
int dep[N],wgt[N];
ll aryf[N],aryg[N<<1],*f[N],*g[N],*fcnt,*gcnt;
ll ans;

void DFS1(ci u,ci fa){
	for(int it=Gr[u];it;it=Gr.nxt[it]){
		int v=Gr.to[it];
		if(v==fa) continue;
		DFS1(v,u);
		if(dep[v]>dep[u]) dep[u]=dep[v],wgt[u]=v;
	}
	dep[u]++;
}
void DFS2(ci u,ci fa){
	if(wgt[u]){
		int v=wgt[u];
		f[v]=f[u]+1,g[v]=g[u]-1;
		DFS2(v,u);
	}
	f[u][0]=1;
	ans+=g[u][0];
	for(int it=Gr[u];it;it=Gr.nxt[it]){
		int v=Gr.to[it];
		if(v==fa || v==wgt[u]) continue;
		alloc(v);
		DFS2(v,u);
		for(int i=0;i<dep[v];i++){
			ans+=g[u][i+1]*f[v][i];
			if(i) ans+=g[v][i]*f[u][i-1];
		}
		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(){
	n=Read();
	for(int i=1;i<n;i++) Gr.AddEdge(Read(),Read());
	DFS1(1,0);
	fcnt=aryf,gcnt=aryg;
	alloc(1);
	DFS2(1,0);
	printf("%lld\n",ans);
	return 0;
}

---

THE END

Thanks for reading!

$$\begin{split} “\ &你—\ 给了我\ 开始和最终\ 放纵的自由与宽容\\ &感谢你借我一半悲伤\ 一半喜悦\ 拼凑\\ &这一秒完整的我\\ &我让我自己梦一场美梦\\ &光与影为彼此停留\\ &我永远为你祈祷此刻\\ &一粒尘埃中的相拥\\ &恨着—\ 爱着—\ 共存着—\ ”\\ ——&\text{《零和Zero-Sum》By 星尘} \end{split}$$

> Linked 零和Zero-Sum