【题解】P5904 [POI2014]HOT-Hotels (长链剖分)

【题解】P5904 [POI2014]HOT-Hotels (长链剖分)

https://www.luogu.com.cn/problem/P5904

题目大意是问一棵树上距离两两相等的三元组的数量

观察到这样的三元组一定有一个交点。一个naive的想法是我们枚举这个交点，然后换根DP就行，分两种情况统计答案(该点是其中恰好三个点的lca，该点是其中恰好两个点的lca)，复杂度\(O(n^2)\)。

但是这样无法通过更大的数据，复杂度是\(O(n^2)\)的主要原因是"该点是其中恰好两个点的lca"不好处理，考虑如何搞一下能够优化复杂度。

实际上不好处理的最主要的原因是，第三个点\(x\)不在我们可以控制的范围内，我们关于他所知的信息只有dis，我们是否可以以\(x\)为交点统计答案，这样所有需要找的那两个点都在\(x\)的范围内了。而且一个好消息是，如果都要找的在子树里，那么这些点的关键信息是深度，可能可以利用长链剖分。

设\(f(i,j)\)表示\(i\)的子树中有多少距离为\(j\)的点，\(g(i,j)\)表示\(dis(\mathrm {lca(a,b)},i)=j\)点对的数量。

设当前点为\(x\),需要转移上来的点是\(t\),转移的情况是：

• \(x\)作为\(\mathrm{lca}(a,b)\)其中\(b\in \mathrm{subtree}(t)\)
• \(f[t],g[t]\)数组分别右移和左移一位加入\(f[x],g[x]\)

数组右移左移一位，如果只有一个儿子的话可以通过指针做到\(O(1)\)，下面的转移可以做到\(O(\min(\mathrm{maxLen}[t],\mathrm{maxLen}[x]))\)，直接长链剖分优化下就万事了。

至于统计答案，也是分两种情况ans+=...就行了

//@winlere
#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>

using namespace std;  typedef long long ll;
inline int qr(){
	int ret=0,f=0,c=getchar();
	while(!isdigit(c))f|=c==45,c=getchar();
	while(isdigit(c)) ret=ret*10+c-48,c=getchar();
	return f?-ret:ret;
}
const int maxn=1e5+5;
ll *f[maxn],*g[maxn],temp[maxn<<2],*p=temp,ans;
vector<int> e[maxn];
void add(int fr,int to){
	e[fr].push_back(to);
	e[to].push_back(fr);
}
int n;
int son[maxn],dep[maxn],maxLen[maxn];

void dfs(int now,int last){
	dep[now]=dep[last]+1;
	for(auto t:e[now])
		if(!dep[t]){
			dfs(t,now);
			if(maxLen[t]+1>maxLen[now])
				maxLen[now]=maxLen[t]+1,son[now]=t;
		}
	if(!maxLen[now]) maxLen[now]=1;
}

void dp(int now){
	if(son[now])
		f[son[now]]=f[now]+1,g[son[now]]=g[now]-1,dp(son[now]);
	f[now][0]=1; ans+=g[now][0];
	for(auto t:e[now])
		if(f[t]==NULL){
			f[t]=p; p+=maxLen[t]<<1;
			g[t]=p; p+=maxLen[t]<<1;
			dp(t);
			for(int j=0;j<maxLen[t];++j)
				ans+=j?f[now][j-1]*g[t][j]:0,ans+=g[now][j+1]*f[t][j];
			for(int j=0;j<maxLen[t];++j){
				g[now][j+1]+=f[now][j+1]*f[t][j];
				if(j) g[now][j-1]+=g[t][j];
				f[now][j+1]+=f[t][j];
			}
		}
}

int main(){
	n=qr();
	for(int t=1;t<n;++t) add(qr(),qr());
	dfs(1,0);
	f[1]=p; p+=maxLen[1]<<1;
	g[1]=p; p+=maxLen[1]<<1;
	dp(1);
	printf("%lld\n",ans);
	return 0;
}