[POI2014] HOT-Hotels 加强版题解

好好好，太好了这题，太好了。

---

首先有一点是很明显的：

对于一个合法的答案 \((i,j,k)\)，必有一点 \(p\)，使 \(dis(i,p)=dis(j,p)=dis(k,p)\) 且三点到 \(p\) 的路径没有交叉。

那所以考虑设 \(g_{u,d}\) 表示 \(u\) 子树内，有多少个二元组 \((i,j)\)，满足 \(dis(i,lca(i,j))=dis(j,lca(i,j))=dis(u,lca(i,j))+d\)。为了辅助运算，设 \(f_{u,d}\) 表示 \(u\) 子树内，有多少个点 \(i\) 满足 \(dis(i,u)=d\)。

对于 \(f_{u,d}\)，转移方程为 \(f_{u,d}=\sum\limits_{v\in son_u}f_{v,d-1}\)，特殊的，\(f_{u,0}=1\)。

对于 \(g_{u,d}\) 来说，有两种情况：

1. 直接从儿子转移而来：\(g_{u,d}+=g_{v,d+1}\)。

2. 两个儿子合并：\(g_{u,d}+=f_{x,d}\times f_{y,d-1}\)

这样实现时间复杂度 \(O(n^2)\)，可以通过原题，但不能通过加强版。

考虑长链剖分优化。假如我们让 \(u\) 点直接继承它的长儿子 \(sn_u\) 的革命果实（由于现在其他孩子还没往里加，所以 \(g\) 的二号公式值为零），暴力合并其他儿子，可以达到单次均摊 \(O(1)\)。

但是为了防止空间爆炸，我们还需要使用指针进行继承。

时空复杂度皆为 \(O(n)\)。

#include<bits/stdc++.h>
#define int long long
#define $ inline
using namespace std;
const int N=1e5+5;
$ int read(){
	char c=getchar();
	int res=0,f=1;
	while(c<'0'||c>'9'){
		if(c=='-')f=-1;
		c=getchar();
	}while(c>='0'&&c<='9')
		res=res*10+c-'0',c=getchar();
	return res*f;
}vector<int>ve[N];int mxp[N];
int n,ans,mp[N],*f[N],*g[N];
int sn[N],tmp[N*4],*id=tmp;
void dfs(int x,int fa){
	for(auto y:ve[x]){
		if(y==fa) continue;
		dfs(y,x);
		if(mxp[x]<mxp[y])
			sn[x]=y,mxp[x]=mxp[y];
	}mxp[x]++;
}void dp(int x,int fa){
	if(!sn[x]) return f[x][0]=1,void();
	f[sn[x]]=f[x]+1,g[sn[x]]=g[x]-1;
	dp(sn[x],x),f[x][0]=1,ans+=g[x][0];
	for(auto y:ve[x]){
		if(y==fa||y==sn[x]) continue;
		f[y]=id,id+=mxp[y]<<1;
		g[y]=id,id+=mxp[y]<<1,dp(y,x);
		for(int j=0;j<mxp[y];j++){
			if(j) ans+=f[x][j-1]*g[y][j];
			ans+=g[x][j+1]*f[y][j];
		}for(int j=0;j<mxp[y];j++){
			g[x][j+1]+=f[x][j+1]*f[y][j];
			if(j) g[x][j-1]+=g[y][j];
			f[x][j+1]+=f[y][j];
		}
	}
}signed main(){
	ios::sync_with_stdio(0);
	cin.tie(0),cout.tie(0);
	n=read();
	for(int i=1,u,v;i<n;i++){
		u=read(),v=read();
		ve[u].push_back(v);
		ve[v].push_back(u);
	}dfs(1,0),f[1]=id,id+=mxp[1]<<1;
	g[1]=id,id+=mxp[1]<<1,dp(1,0);
	cout<<ans;
	return 0;
}