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

传送门
代码：
长链剖分好题。
题意：给你一棵树，问树上选三个互不相同的节点，使得这个三个点两两之间距离相等的方案数。

---

思路：
先考虑$dp$。
$f_{i,j}$表示$i$子树中离$i$距离为$j$的点数，$g_{i,j}$表示$i$子树中所有满足$dist(lca(u,v),i)-dist(lca(u,v),i)=j$的点对数。
不难发现这两个状态是可以进行拼接的。
因此对于一对父亲和孩子$(p,v)$有如下转移式：
$ans+=g_{v,i+1}\ast f_{p,i}+g_{p,i}\ast f_{v,i-1}$
$g_{p,i}+=g_{v,i-1}+f_{v,i-1}\ast f_{p,i}$
$f_{p,i}+=f_{v,i-1}$
然后用长链剖分优化一下即可。
代码：

#include
#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;
}