[FROM WOJ]#3604 [POI2014]HOT_Hotels(数据有加强）

#3604 [POI2014]HOT_Hotels(数据有加强）

题面
有一个树形结构的宾馆，n (1≤n≤100 000)个房间，n-1条无向边，每条边的长度相同，任意两个房间可以相互到达。吉丽要给他的三个妹子各开（一个）房（间）。三个妹子住的房间要互不相同（否则要打起来了），为了让吉丽满意，你需要让三个房间两两距离相同。 有多少种方案能让吉丽满意？

输入
第一行一个数n。
接下来n-1行，每行两个数x,y，表示x和y之间有一条边相连。

输出
让吉丽满意的方案数。

样例输入
7
1 2
5 7
2 5
2 3
5 6
4 5

样例输出
5

SOL
设：
f(i,j)=∑x∈subtree(i) [dis(i,x)=j]
g(i,j)=∑lca(x,y)∈subtree(i) [dis(x,lca)=dis(y,lca)=dis(i,lca)+j]
每次将一个儿子y转移到x：
ans+=∑f(y,i)∗g(x,i+1)+g(y,i)∗f(x,i−1)
f(x,i+1)+=f(y,i)
g(x,i−1)+=g(y,i)
g(x,i+1)+=f(x,i+1)∗f(y,i) (lca为x)

代码：

#include<bits/stdc++.h>
using namespace std;
#define re register
#define int long long
inline int rd(){
	int data=0;static char ch=0;ch=getchar();
	while(!isdigit(ch))ch=getchar();
	while(isdigit(ch))data=(data<<1)+(data<<3)+(ch^48),ch=getchar();
	return data;
}
inline void write(int x){if(x>9)write(x/10);putchar(x%10+'0');}
const int N=1e6+5;
bool vis[N];
int f[N*10],g[N*10],n,ans,li,cnt,first[N],tcnt,pos[N],dfn[N],top[N],len[N],lson[N],fa[N],dep[N];
struct node{int v,nxt;}e[N<<1];
inline void add(int u,int v){e[++cnt]=(node){v,first[u]};first[u]=cnt;}
inline void dfs1(int u){
	len[u]=1,vis[u]=true;
	for(int re i=first[u];i;i=e[i].nxt){
		int re v=e[i].v;
		if(vis[v])continue;
		fa[v]=u,dep[v]=dep[u]+1,dfs1(v);
		len[u]=max(len[u],len[v]+1);
		if(len[lson[u]]<len[v]+1)lson[u]=v;
	}
}
inline void dfs2(int u){
	dfn[u]=++tcnt;
	if(lson[u])top[lson[u]]=top[u],dfs2(lson[u]);
	for(int re i=first[u];i;i=e[i].nxt){
		int re v=e[i].v;
		if(v==lson[u]||v==fa[u])continue;
		top[v]=v,pos[v]=li+len[v],li+=len[v]<<1,dfs2(v);
	}
}
inline int get_f(int u,int i){return dfn[u]+i;}
inline int get_g(int u,int i){return pos[top[u]]-dep[u]+dep[top[u]]+i;}
inline void cal(int x,int u){
	for(int re i=1;i<=len[u];i++)
		ans+=f[get_f(u,i-1)]*g[get_g(x,i)];
	for(int re i=1;i<=len[u];i++)
		ans+=g[get_g(u,i)]*f[get_f(x,i-1)];
	for(int re i=0;i<=len[u]-1;i++)
		g[get_g(x,i)]+=g[get_g(u,i+1)];
	for(int re i=1;i<=len[u];i++)
		g[get_g(x,i)]+=f[get_f(x,i)]*f[get_f(u,i-1)],f[get_f(x,i)]+=f[get_f(u,i-1)];
}
inline void sol(int u){
	if(!lson[u]){f[get_f(u,0)]=1,g[get_g(u,0)]=0;return;}
	sol(lson[u]);
	g[get_g(u,len[u])]=g[get_g(u,len[u]-1)]=0;
	ans+=g[get_g(u,0)],f[get_f(u,0)]=1;
	for(int re i=first[u];i;i=e[i].nxt){
		int re v=e[i].v;
		if(v==fa[u]||v==lson[u])continue;
		sol(v),cal(u,v);
	}
}
signed main(){
//	freopen("hothotels.in","r",stdin);
//	freopen("hothotels.out","w",stdout);
	n=rd();
	for(int re i=1;i<n;i++){int re x=rd(),y=rd();add(x,y),add(y,x);}
	dep[1]=1,dfs1(1),top[1]=1,pos[1]=len[1],li=len[1]<<1,dfs2(1),sol(1);
	write(ans);
	return 0;
}