【BZOJ4543】[POI2014] Hotel加强版（长链剖分优化DP）

点此看题面

• 给定一棵\(n\)个点的树。
• 问有多少种在树上选三个点的方案，满足这三个点两两距离相等。
• \(n\le 10^5\)

暴力换根

先考虑原题【BZOJ3522】[POI2014] Hotel，与此题区别只在于\(n\le5\times10^3\)。

容易发现，符合条件的三个点之间路径的交点必然满足到三点距离相等。

于是我们枚举交点作为根，只要求出有多少种方案在三个不同的子树中各选出一点，满足这三个点到根距离相等。

这应该是非常简单的。

不换根的暴力

现在我们必须固定根不动了。

设\(f_{x,i},g_{x,i}\)，分别表示\(x\)子树中到\(x\)距离为\(i\)的点数和\(x\)子树中选出两点能和一个到\(x\)距离为\(i\)的点匹配的方案数。

考虑每次枚举一个子节点\(y\)更新\(x\)的信息，得到下列式子：

\[ans\texttt{+=}\sum (f_{x,i-1}\times g_{y,i}+g_{x,i+1}\times f_{y,i})\\ g_{x,i+1}\texttt{+=}f_{x,i+1}\times f_{y,i}\\ f_{x,i+1}\texttt{+=}f_{x,i}\\ g_{x,i-1}\texttt{+=}g_{x,i} \]

注意这里的转移有些是有一定顺序的，不过结合定义应该不难理解。

长链剖分优化

这种东西应该是典型的长链剖分优化\(DP\)的式子。

要注意的是，\(g_{x,i}\)每次会从\(g_{son,i+1}\)处转移得到，与一般长链剖分起始位置不断左移不同，是一个右移的过程，因此对于它的内存池要开成两倍。

代码：\(O(n)\)

#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Reg register
#define RI Reg int
#define Con const
#define CI Con int&
#define I inline
#define W while
#define N 100000
#define LL long long
#define add(x,y) (e[++ee].nxt=lnk[x],e[lnk[x]=ee].to=y)
using namespace std;
int n,ee,lnk[N+5];struct edge {int to,nxt;}e[N<<1];
int s[N+5],l[N+5];I void dfs1(CI x,CI lst=0)//第一遍dfs，预处理长儿子
{
	for(RI i=lnk[x];i;i=e[i].nxt) e[i].to^lst&&
		(dfs1(e[i].to,x),l[e[i].to]>l[s[x]]&&(s[x]=e[i].to));l[x]=l[s[x]]+1;
}
int t_f,t_g,id_f[N+5],id_g[N+5];I void dfs2(CI x,CI p1,CI p2,CI lst=0)//第二遍dfs，分配内存
{
	id_f[x]=p1,id_g[x]=p2,s[x]&&(dfs2(s[x],p1+1,p2-1,x),0);
	for(RI i=lnk[x];i;i=e[i].nxt) e[i].to^lst&&e[i].to^s[x]&&
		(t_f+=l[e[i].to],t_g+=l[e[i].to]<<1,dfs2(e[i].to,t_f-l[e[i].to]+1,t_g-l[e[i].to]+1,x),0);
}
#define f(x,i) F[id_f[x]+i]
#define g(x,i) G[id_g[x]+i]
LL ans,F[N+5],G[2*N+5];I void DP(CI x,CI lst=0)//动态规划
{
	RI i,j;for(s[x]&&(DP(s[x],x),0),f(x,0)=1,ans+=g(x,0),i=lnk[x];i;i=e[i].nxt)//注意与长儿子的答案
	{
		if(e[i].to==lst||e[i].to==s[x]) continue;DP(e[i].to,x);
		for(j=0;j^l[e[i].to];++j) ans+=g(x,j+1)*f(e[i].to,j);//计算答案
		for(j=1;j^l[e[i].to];++j) ans+=f(x,j-1)*g(e[i].to,j);//计算答案
		for(j=0;j^l[e[i].to];++j) g(x,j+1)+=f(x,j+1)*f(e[i].to,j);//凑成一对
		for(j=0;j^l[e[i].to];++j) f(x,j+1)+=f(e[i].to,j);//直接转移
		for(j=1;j^l[e[i].to];++j) g(x,j-1)+=g(e[i].to,j);//直接转移
	}
}
int main()
{
	RI i,x,y;for(scanf("%d",&n),i=1;i^n;++i) scanf("%d%d",&x,&y),add(x,y),add(y,x);
	return dfs1(1),t_f=l[1],t_g=l[1]<<1,dfs2(1,1,l[1]+1),DP(1),printf("%lld\n",ans),0;
}