DP优化——长剖优化DP

长链剖分

在长链剖分中重儿子的定义为：子树深度最大的儿子。
其余就和重剖一样了。
下面是核心代码：

void dfs(int u){
	mxdep[u]=0;  //子树内最大深度
	for(int i=head[u];i;i=Next[i]){
		int v=to[i];
		if(v==fa[u]) continue;
		fa[v]=u;
		dfs(v);
		if(mxdep[v]>=mxdep[son[u]]) son[u]=v;
	}
	if(son[u]) mxdep[u]=mxdep[son[u]]+1;
}

应用

长链剖分的应用很多，但因为我们放在了 DP 优化合集里，所以我们只介绍他如何优化 DP。

要用他优化 DP 我们只需要知道他的一个性质就够了：所有长链的长度之和为 O(n)。

能用长链剖分优化的 DP 满足其中一维 dp 状态跟深度有关，此时如果我们对每一个点那一维都开一个 \(O(n)\) 的数组，那么空间和时间起码就是 \(O(n^2)\) 的了。
但其实对于每个点 \(u\) 他那一维的大小实际上只需要开 \(maxdep_u\) 的大小即可，有很多是浪费的。
所以我们借鉴 dsu on tree 的思想，每个点直接继承重儿子的 dp 数组，然后暴力合并轻儿子的 dp 数组。
此时每一条长链只会在链顶向其父亲转移时会被暴力合并，所以总复杂度就是所有长链的长度之和 \(O(n)\)。

在实现继承数组这一操作时，我们往往使用指针，在遍历到链顶时提前为整条长链开好空间，然后这条长链上的点公用这段空间。
根据 dp 转移的不同可能会需要用不同的写法。

下面以一道例题来详细介绍。

例题

几乎所有介绍长剖优化 DP 的博客（包括 oi-wiki） 都以 CF1009F 作为例题，所以本博客这里不再以他为例题介绍。

[POI2014] HOT-Hotels 加强版

画个图容易发现满足条件的三个点 \((x,y,z)\) 一定满足两者之一：

1. 他们的 \(LCA\) 是同一个，且他们距离 \(LCA\) 的距离都为 \(d\)。
2. 他们其中两个（不妨设为 \(x,y\)）距离他们的 \(LCA\)（下面记为 \(u\)） 距离都为 \(d\)，并且 \(z\) 和 \(u\) 的距离也为 \(d\)。而且 \(LCA(x,z)=LCA(y,z)=lCA(u,z)=v\)，即要保证他们之间的距离为 \(2d\)。

其实第 \(1\) 种情况是第 \(2\) 种情况的特殊情况。

为了避免算重我们把每个三元组算在最高处，比如在第一种情况里就算在 \(LCA\) 处，在第二种情况里就算在 \(v\) 处。
然后进行 dp：
\(f_{i,j}\) 表示和 \(i\) 子树内和 \(i\) 距离为 \(j\) 的点的个数。
\(g_{i,j}\) 表示 \(i\) 子树内的点对 \((x,y)\) 满足 \(LCA(x,y)=lca，dist(x,lca)=dist(y,lca)=d\)，且 \(dist(i,lca)=d-j\) 的 \((x,y)\) 的个数。

转移考虑树形背包，假设现在合并进来 \(u\) 的一个子树 \(v\)：
\(f\) 的转移是简单的（和 CF1009F 一样）：\(f_{u,j+1} += f_{v,j}\)。
对于 \(g\) 的转移分两种：

1. 点对来自 \(v\) 子树内：\(g_{u,j-1} += g_{v,j}\)
2. 点对分别来自前面的子树和 \(v\) 子树，此时他们的 \(lca\) 就是 \(u\)：\(g_{u,j+1} += f_{u,j+1}\times f_{v,j}\)。

然后考虑对 \(ans\) 的贡献：

1. 最后加上 \(g_{u,0}\)：此时 \(u\) 就是那第三个点。
2. 还是考虑在合并进 \(v\) 子树的时候，此时的贡献分为 \(v\) 子树提供单点和 \(v\) 子树提供点对：\(g_{u,j+1} \times f_{v,j} + f_{u,j-1} \times g_{v,j}\)。

但其实当 \(j=1\) 时，\(f_{u,j-1}\times g_{v,j}\) 就是第 \(1\) 种情况的贡献。
所以不需要额外把第 \(1\) 种情况算进去。

第二维跟子树最大深度同阶，可以长剖优化成 \(O(n)\)。
要注意的是 \(g\) 数组在继承重儿子的时候是 \(g_{u,j-1}=g_{son_u,j}\)，而不是像 \(f\) 那样 \(f_{u,j+1}=f_{son_u,j}\)，所以内存分配的方式有所不同。

code

#include<bits/stdc++.h>
#define int long long  
using namespace std;
const int N=1e5+5;

inline int read(){
	int w=1,s=0;
	char c=getchar();
	for(;c<'0'||c>'9';w*=(c=='-')?-1:1,c=getchar());
	for(;c>='0'&&c<='9';s=s*10+c-'0',c=getchar());
	return w*s;
}
int n,tot,head[N],to[N<<1],Next[N<<1];
void add(int u,int v){
	to[++tot]=v,Next[tot]=head[u],head[u]=tot;
}
int mxdep[N],son[N],fa[N],*f[N],*g[N],ans,tmp1[N<<2],tmp2[N<<2],*id1=tmp1,*id2=tmp2;  //f,g 只是指向内存开头的指针，tmp1,tmp2 是用来分配内存的数组
void dfs(int u){
	mxdep[u]=0;
	for(int i=head[u];i;i=Next[i]){
		int v=to[i];
		if(v==fa[u]) continue;
		fa[v]=u;
		dfs(v);
		if(mxdep[v]>=mxdep[son[u]]) son[u]=v;
	}
	if(son[u]) mxdep[u]=mxdep[son[u]]+1;
}
void dfs1(int u){
	if(son[u]){
		f[son[u]]=f[u]+1;  // 此时对于 son[u] 来讲的 dp 数组的第 i 位，就是对于 u 来讲的 dp 数组的第 i+1 位，下同
		g[son[u]]=g[u]-1;
		dfs1(son[u]);
		ans += g[son[u]][1];
	}
	f[u][0]=1;
	for(int i=head[u];i;i=Next[i]){
		int v=to[i];
		if(v==fa[u] || v==son[u]) continue;
		f[v]=id1; id1+=mxdep[v]+1;  //在链顶分配内存，对于 f 只需要往后分配 maxdep[v] 的长度即可
		id2+=mxdep[v]; g[v]=id2; id2+=mxdep[v]+1;   //对于 g 不仅需要往后分配还需要往前分配（因为重儿子 dp 数组的开头在他的开头前面）
		dfs1(v);
		for(int j=0;j<=mxdep[v];j++) ans += g[u][j+1] * f[v][j] + ((j>0) ? (f[u][j-1] * g[v][j]) : 0); 
		for(int j=0;j<=mxdep[v];j++){
			g[u][j+1]+=f[u][j+1]*f[v][j];
			if(j>0) g[u][j-1]+=g[v][j];
		}
		for(int j=0;j<=mxdep[v];j++) f[u][j+1]+=f[v][j];
	}
}
signed main(){
	n=read();	
	for(int i=1;i<n;i++){
		int u=read(),v=read();
		add(u,v),add(v,u); 
	}
	
	dfs(1);
	f[1]=id1; id1+=mxdep[1]+1;
	id2+=mxdep[1]; g[1]=id2; id2+=mxdep[1]+1; 
	dfs1(1);
	
	printf("%lld\n",ans);
	return 0;
}