2019.01.08 bzoj4543: [POI2014]Hotel加强版（长链剖分+dp）

传送门

代码：

长链剖分好题。

题意：给你一棵树，问树上选三个互不相同的节点，使得这个三个点两两之间距离相等的方案数。

---

思路：

先考虑dpdpdp。

fi,jf_{i,j}fi,j​表示iii子树中离iii距离为jjj的点数，gi,jg_{i,j}gi,j​表示iii子树中所有满足dist(lca(u,v),i)−dist(lca(u,v),i)=jdist(lca(u,v),i)-dist(lca(u,v),i)=jdist(lca(u,v),i)−dist(lca(u,v),i)=j的点对数。

不难发现这两个状态是可以进行拼接的。

因此对于一对父亲和孩子(p,v)(p,v)(p,v)有如下转移式：

ans+=gv,i+1∗fp,i+gp,i∗fv,i−1ans+=g_{v,i+1}*f_{p,i}+g_{p,i}*f_{v,i-1}ans+=gv,i+1​∗fp,i​+gp,i​∗fv,i−1​

gp,i+=gv,i−1+fv,i−1∗fp,ig_{p,i}+=g_{v,i-1}+f_{v,i-1}*f_{p,i}gp,i​+=gv,i−1​+fv,i−1​∗fp,i​

fp,i+=fv,i−1f_{p,i}+=f_{v,i-1}fp,i​+=fv,i−1​

然后用长链剖分优化一下即可。

代码：

#include<bits/stdc++.h>
#define ri register int
using namespace std;
inline int read(){
	int ans=0;
	char ch=getchar();
	while(!isdigit(ch))ch=getchar();
	while(isdigit(ch))ans=(ans<<3)+(ans<<1)+(ch^48),ch=getchar();
	return ans;
}
typedef long long ll;
const int N=1e5+5;
vector<int>e[N];
int n,hson[N],dep[N],mdep[N],len[N],fa[N];
ll tmp[N<<2],*f[N],*g[N],*now=tmp,ans=0;
void dfs1(int p){
	for(ri i=0,v;i<e[p].size();++i){
		if((v=e[p][i])==fa[p])continue;
		fa[v]=p,dep[v]=mdep[v]=dep[p]+1,dfs1(v),mdep[p]=max(mdep[p],mdep[v]);
		if(mdep[v]>mdep[hson[p]])hson[p]=v;
	}
	len[p]=mdep[p]-dep[p]+1;
}
void dfs2(int p){
	if(hson[p])f[hson[p]]=f[p]+1,g[hson[p]]=g[p]-1,dfs2(hson[p]);
	f[p][0]=1,ans+=g[p][0];
	for(ri i=0,v;i<e[p].size();++i){
		if((v=e[p][i])==fa[p]||v==hson[p])continue;
		f[v]=now,now+=2*len[v],g[v]=now,now+=2*len[v],dfs2(v);
		for(ri j=0;j<len[v];++j){
			ans+=f[v][j]*g[p][j+1];
			if(j)ans+=g[v][j]*f[p][j-1];
		}
		for(ri j=0;j<len[v];++j){
			g[p][j+1]+=f[v][j]*f[p][j+1];
			if(j)g[p][j-1]+=g[v][j];
			f[p][j+1]+=f[v][j];
		}
	}
}
int main(){
	freopen("lx.in","r",stdin);
	while(n=read(),n){
		ans=0;
		memset(tmp,0,sizeof(tmp));
		memset(hson,0,sizeof(hson));
		memset(mdep,0,sizeof(mdep));
		for(ri i=1;i<=n;++i)e[i].clear();
		for(ri i=1,u,v;i<n;++i)u=read(),v=read(),e[u].push_back(v),e[v].push_back(u);
		dfs1(1),now=tmp,f[1]=now,now+=2*len[1],g[1]=now,now+=2*len[1],dfs2(1),cout<<ans<<'\n';
	}
	return 0;
}