【UOJ 33】树上GCD（启发式合并）

传送门

题意

给出一棵 \(n\) 个结点的有根树。记 \(d_u\) 为点 \(u\) 的深度。对于树上两点 \(u, v(u \neq v)\)，定义 \(w(u, v)\) 为 \(\gcd(d_u - d_{\operatorname{lca}(u, v)}, d_v - d_{\operatorname{lca}(u, v)})\)。

对于每个 \(1 \le i \le n-1\)，求出使得 \(w(u, v) = i\) 的点对个数。

\(n \le 2\times10^5\)。

分析

首先只要算出 \(a_i\) 表示使得 \(w(u, v)\) 为 \(i\) 倍数的点对个数，最后对 \(a\) 做一遍容斥就能得到答案数组。

设阈值 \(m\)，分别考虑不超过 \(m\) 和大于 \(m\) 的 \(i\)。

我们可以线性求 \(a_i\) 单项。设 \(f_u\) 表示 \(u\) 子树内深度与 \(d_u\) 在模 \(i\) 意义下同余的点数，\(g_u\) 表示 \(u\) 子树内深度与 \(d_{\operatorname{fa}_u}\) 同余的点数。只要预处理出每个点的 \(i\) 级祖先，就可以按照 dfs 序倒序转移求出这两个数组，之后就能算答案了。用这个方法做 \(i \le m\) 的部分，时间复杂度为 \(O(nm)\)。

对于 \(i > m\) 的部分，设 \(c_{u, l}\) 表示 \(u\) 子树内距离 \(u\) 为 \(l\) 的点数。我们把 \(c_u\) 看作一个数组，先考虑怎么求它。我们找到 \(v\) 表示 \(u\) 的儿子中子树深度最大的，令 \(c_u\) 的初值为 \(c_v\) 前添一项 \(1\)。对于其他儿子，暴力把它们的数组加上来即可。怎么在数组前添一项呢？我们用 vector 维护翻转过的数组，只要用 push_back 即可。这样做的时间复杂度是啥呢？是线性，因为每次暴力加的操作都可以均摊倒一个结点上。

事实上这就是长链剖分。一种更好写的做法是启发式合并，我们按照 dfs 序倒序，每次把一个点的数组合并到父亲上。要确保是把较小的数组暴力加到较大的数组上。因此我们要支持交换两个 vector，可以用 V1.swap(V2)，复杂度是常数。分析一下就会发现这个做法的时间复杂度和长链剖分一样，证明也类似。

现在考虑怎么求答案。我们在合并时枚举 \(i\)，每次跳 \(i\) 步地遍历两个数组，将各自的总和相乘贡献到 \(a_i\)。这样的瓶颈在于遍历较长的数组上，但是由于 \(i\) 大于 \(m\)，跳的次数不会超过 \(\frac{n}{m}\)。而 \(i\) 不超过较小数组的长度，因此枚举量可以均摊到每个结点上，是线性的。因此这部分的时间复杂度为 \(O(\frac{n^2}{m})\)。

令 \(m = \sqrt n\)，整个问题即可做到 \(O(n \sqrt n)\)。实际上第二部分常数很小，因此 \(m\) 较小时跑得快。

实现

#include<bits/stdc++.h>
using namespace std;
#define rep(i,a,b) for(int i=a;i<=b;i++)
#define per(i,a,b) for(int i=a;i>=b;i--)
#define rep2(i,a,b,c) for(int i=a;i<=b;i+=c)
#define pb push_back
typedef long long ll;
const int N=2e5+5,m=20;
int n,p[N],d[N],ex[N],pk[N],f[N],g[N],x,y;
ll ans[N];
vector<int>A[N];
int main(){
	scanf("%d",&n);
	rep(i,2,n)scanf("%d",&p[i]),ex[d[i]=d[p[i]]+1]++,pk[i]=i;
	per(i,n,1)ex[i]+=ex[i+1];
	rep(k,1,m){
		per(i,n,2)g[pk[i]]+=++f[i],ans[k]+=ll(g[i])*f[p[i]],f[p[i]]+=g[i];
		rep(i,1,n)pk[i]=p[pk[i]],f[i]=g[i]=0;
	}
	per(i,n,2){A[i].pb(1);
		int l=A[i].size(),r=A[p[i]].size();
		if(l>r)A[i].swap(A[p[i]]),swap(l,r);
		l=A[i].size(),r=A[p[i]].size();
		rep(k,m+1,l){
			rep2(j,k,l,k)x+=A[i][l-j];
			rep2(j,k,r,k)y+=A[p[i]][r-j];
			ans[k]+=ll(x)*y,x=y=0;
		}
		rep(j,1,l)A[p[i]][r-j]+=A[i][l-j];
	}
	per(i,n,1)rep2(j,i*2,n,i)ans[i]-=ans[j];
	rep(i,1,n-1)printf("%lld\n",ans[i]+ex[i]);
	return 0;
}