bzoj 4543: [POI2014]Hotel加强版

Description

给出一棵树求三元组 \((x,y,z)\,,x<y<z\) 满足三个点两两之间距离相等,求三元组的数量

Solution

考虑暴力 \(DP\)
设 \(f[i][j]\) 表示距离点 \(i\) 的子树内距离为 \(j\) 的点的数量
设 \(g[i][j]\) 表示 \(i\) 子树内的一个二元组 \((x,y)\) 满足 \(dis(x,lca)=dis(y,lca)=dis(i,lca)+j\) 的二元组的数量,可以视为等待与子树外合并的二元组的数量
显然有转移:
\(ans+=g[x][0]\)
\(ans+=f[son][j]*g[x][j+1]+f[x][j-1]*g[son][j]\)

\(f[x][j]=f[son][j-1]\)
\(g[x][j]=g[son][j+1]\)
\(g[x][j]=f[x][j]*f[son][j-1]\)

这样转移是 \(O(n^2)\) 的
对于深度为下标的树形 \(DP\),考虑长链剖分优化:
在 \(DP\) 转移之前, \(f[x],g[x]\) 是没有值的,初值要设为某个儿子的 \(DP\) 值
并且这个时候的转移仅仅是 \(f[x][j]=f[son][j-1]\),\(g[x][j]=g[son][j+1]\)
相当于一个数组位移,直接用指针优化,可以做到 \(O(1)\),所以用所在链最长的儿子来赋初值复杂度最优,其余儿子暴力转移
这样做复杂度均摊就是 \(O(n)\) 的了,一条链只会在链顶被枚举到,且链不相交,所以每个点只会被枚一次
空间对于每一个链动态开空间,空间复杂度也是 \(O(n)\) 的

#include<bits/stdc++.h>
using namespace std;
template<class T>void gi(T &x){
	int f;char c;
	for(f=1,c=getchar();c<'0'||c>'9';c=getchar())if(c=='-')f=-1;
	for(x=0;c<='9'&&c>='0';c=getchar())x=x*10+(c&15);x*=f;
}
typedef long long ll;
const int N=1e5+10;
int n,head[N],nxt[N*2],to[N*2],num=0,dep[N],mx[N];
inline void link(int x,int y){nxt[++num]=head[x];to[num]=y;head[x]=num;}
ll lis[N*5],*f[N],*g[N],*st=lis+5,ans=0;
inline void dfs(int x,int fa){
	mx[x]=x;
	for(int i=head[x],u;i;i=nxt[i]){
		if((u=to[i])==fa)continue;
		dep[u]=dep[x]+1;dfs(u,x);
		if(dep[mx[u]]>dep[mx[x]])mx[x]=mx[u];
	}
	for(int i=head[x],u;i;i=nxt[i]){
		if((u=to[i])==fa || (mx[u]==mx[x] && x!=1))continue;
		st+=dep[mx[u]]-dep[x]+1;
		f[mx[u]]=st;
		g[mx[u]]=++st;
		st+=(dep[mx[u]]-dep[x])*2+1;
	}
}
inline void dfs2(int x,int fa){
	for(int i=head[x],u;i;i=nxt[i]){
		if((u=to[i])==fa)continue;
		dfs2(u,x);
		if(mx[u]==mx[x])f[x]=f[u]-1,g[x]=g[u]+1;
	}
	f[x][0]=1;ans+=g[x][0];
	for(int i=head[x],u;i;i=nxt[i]){
		if((u=to[i])==fa || mx[u]==mx[x])continue;
		for(int j=dep[mx[u]]-dep[x];j>=0;j--)
			ans+=g[x][j+1]*f[u][j]+g[u][j]*(j?f[x][j-1]:0);
		for(int j=dep[mx[u]]-dep[x];j>=0;j--){
			f[x][j+1]+=f[u][j];
			g[x][j]+=g[u][j+1];
			g[x][j]+=f[x][j]*(j?f[u][j-1]:0);
		}
	}
}
int main(){
  freopen("pp.in","r",stdin);
  freopen("pp.out","w",stdout);
  cin>>n;
  int x,y;
  for(int i=1;i<n;i++)gi(x),gi(y),link(x,y),link(y,x);
  dfs(1,1);dfs2(1,1);
  cout<<ans<<endl;
  return 0;
}