Codeforces 434E - Furukawa Nagisa's Tree（三元环+点分治）

Codeforces 题面传送门 & 洛谷题面传送门

场号 hopping，刚好是我的学号（指 round 的编号）

注：下文中分别用 \(X,Y,K\) 代替题目中的 \(x,y,k\)

注意到这东西长得有点像三元环，因此考虑往三元环方面考虑。我们重新建立一张图，对于两点 \(x,y\)，如果 \(x\to y\) 路径对应的权值为 \(X\)（即 \(x\to y\) 这条路径属于 Furukawa Nagisa），那么我们就在 \(x\to y\) 之间连一条权值为 \(1\) 的边，否则我们在 \(x\to y\) 之间连一条权值为 \(0\) 的边，那么我们要求的即为有多少个三元组 \((x,y,z)\)，满足 \(x\to y,y\to z,x\to z\) 之间全是权值为 \(1\) 的边，或者全是权值为 \(0\) 的边。我们考虑借鉴竞赛图三元环计数的一个小 trick：从反面考虑。具体来说我们考虑不合法的情况长什么样，不难发现，对于每种不合法的情况都恰好有两个点，满足与它相连的两条边恰好一条权值为 \(0\)，一条权值为 \(1\)。我们记这样的点为这个三元组的“代表点”。我们考虑对每个点分别统计有多少个三元组，满足其在该三元组中作为一个代表点出现，那么一个不合法的三元组会恰好被计算两次，因此将这个数目除以 \(2\) 就是不合法的三元组个数。

那么怎么计算每个点在多少个三元组中作为代表点呢？注意到对于一个点 \(x\) 而言，与其相连的点可以被分为四类：\(y\to x\)，权值为 \(0/1\)，\(x\to y\)，权值为 \(0/1\)。因此考虑将与 \(x\) 相连的两条边，一条权值为 \(0\)，一条权值为 \(1\) 的情况分为以下四类：

• \(y\to x\)​，权值为 \(0\)​，\(z\to x\)​，权值为 \(1\)​
• \(y\to x\)，权值为 \(0\)，\(x\to z\)，权值为 \(1\)
• \(x\to y\)，权值为 \(0\)，\(z\to x\)，权值为 \(1\)
• \(x\to y\)，权值为 \(0\)，\(x\to z\)，权值为 \(1\)

不难发现对于第一种和第四种情况，\(y,z\) 之间的边既可以是 \(y\to z\)，也可以是 \(z\to x\)，因此贡献应该乘 \(2\)，而对于第二种和第三种情况，\(y,z\) 之间边的指向唯一。故如果我们已经求出 \(c1_x\) 表示有多少个 \(y\) 满足 \(y\to x\) 的边权值为 \(1\)，\(c2_x\) 表示有多少个 \(y\) 满足 \(x\to y\) 的边的权值为 \(1\)，那么一个点 \(x\) 的贡献即可写作：

\[2c1_x·(n-c1_x)+2c2_x·(n-c2_x)+c1_x·(n-c2_x)+c2_x·(n-c1_x)
\]

接下来考虑怎样求出 \(c1_x,c2_x\)​，注意到这里涉及树上路径统计，因此考虑点分治，于是问题可以转化为如何求出经过当前分治中心 \(x\)​ 的路径的贡献。对于 \(x\)​ 所在的连通块中的点 \(y\)​，我们考虑一遍 DFS 找出 \(x\to y\)​ 路径上所有点 \(x=p_1,p_2,p_3,\cdots,p_k=y\)​，那么我们考虑设 \(s_y=\sum\limits_{i=2}^ka_{p_i}·K^{k-i},t_y=\sum\limits_{i=1}^ka_{p_i}·K^{i-1}\)​，那么对于一条 \(u\to v\)​ 的路径，其权值可以写作 \(s_u+t_v·K^{dep_u}\)​。我们考虑从左到右遍历所有子树，那么 \(c1_v\)​ 会加上满足 \(s_u+t_v·K^{dep_u}\equiv X\pmod{Y}\)​ 的 \(u\)​ 的个数，这个可以开一个桶 \(b1\)​，然后每加入一个点 \(u\)​ 时，就在 \((X-s_u)·K^{-dep_u}\)​ 位置加 \(1\)​，然后通过调用 \(b1_{t_v}\)​ 即可求出符合要求的 \(u\) 的个数。\(c2_v\) 的贡献也同理，一个点 \(u\) 会对 \(c2_v\) 产生贡献当且仅当 \(s_v+t_u·K^{dep_v}\equiv X\pmod{Y}\)，化简可得 \(t_u\equiv(X-s_v)·K^{-dep_u}\pmod{Y}\)，同样开个桶 \(b2\) 维护一下即可。

时间复杂度 \(n\log^2n\)

const int MAXN=1e5;
const int INF=0x3f3f3f3f;
int n,X,Y,k,a[MAXN+5];
int hd[MAXN+5],to[MAXN*2+5],nxt[MAXN*2+5],ec=0;
int qpow(int x,int e){
	int ret=1;
	for(;e;e>>=1,x=1ll*x*x%Y) if(e&1) ret=1ll*ret*x%Y;
	return ret;
}
int pw[MAXN+5],ipw[MAXN+5];
void adde(int u,int v){to[++ec]=v;nxt[ec]=hd[u];hd[u]=ec;}
int siz[MAXN+5],cent=0,mx[MAXN+5];bool vis[MAXN+5];
void findcent(int x,int f,int tot){
	siz[x]=1;mx[x]=0;
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==f||vis[y]) continue;
		findcent(y,x,tot);siz[x]+=siz[y];
		chkmax(mx[x],siz[y]);
	} chkmax(mx[x],tot-siz[x]);
	if(mx[x]<mx[cent]) cent=x;
}
int dep[MAXN+5],f1[MAXN+5],f2[MAXN+5];
void getdep(int x,int f){
//	printf("%d %d %d %d\n",x,dep[x],f1[x],f2[x]);
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==f||vis[y]) continue;
		dep[y]=dep[x]+1;f1[y]=(1ll*f1[x]*k+a[y])%Y;
		f2[y]=(f2[x]+1ll*a[y]*pw[dep[y]])%Y;
		getdep(y,x);
	}
}
vector<int> pt;
int res1[MAXN+5],res2[MAXN+5];
void getpts(int x,int f){
	pt.pb(x);
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==f||vis[y]) continue;
		getpts(y,x);
	}
}
void divcent(int x){
//	printf("divcent %d\n",x);
	vis[x]=1;f1[x]=0;f2[x]=a[x];dep[x]=0;
	map<int,int> cnt1,cnt2;
	#define insert1(x) cnt1[1ll*(X-f1[x]+Y)*ipw[dep[x]]%Y]++
	#define insert2(x) cnt2[f2[x]]++
	#define calc(x) (res1[x]+=cnt1[f2[x]],res2[x]+=cnt2[1ll*(X-f1[x]+Y)*ipw[dep[x]]%Y])
	insert1(x);insert2(x);stack<int> stk;
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(vis[y]) continue;
		f1[y]=a[y];f2[y]=(1ll*a[y]*k+a[x])%Y;
		dep[y]=1;getdep(y,x);
	}
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(vis[y]) continue;
//		printf("y=%d\n",y);
		pt.clear();getpts(y,x);stk.push(y);
		for(int z:pt) calc(z);
		for(int z:pt) insert1(z),insert2(z);
	} cnt1.clear();cnt2.clear();
	while(!stk.empty()){
		int y=stk.top();stk.pop();
//		printf("y=%d\n",y);
		pt.clear();getpts(y,x);
		for(int z:pt) calc(z);
		for(int z:pt) insert1(z),insert2(z);
	} calc(x);
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(vis[y]) continue;
		cent=0;findcent(y,x,siz[y]);divcent(cent);
	}
}
int main(){
	scanf("%d%d%d%d",&n,&Y,&k,&X);
	for(int i=1;i<=n;i++) scanf("%d",&a[i]);
	for(int i=1,u,v;i<n;i++) scanf("%d%d",&u,&v),adde(u,v),adde(v,u);
	pw[0]=ipw[0]=1;pw[1]=k;ipw[1]=qpow(k,Y-2);
	for(int i=2;i<=n;i++){
		pw[i]=1ll*pw[i-1]*pw[1]%Y;
		ipw[i]=1ll*ipw[i-1]*ipw[1]%Y;
	} mx[0]=INF;findcent(1,0,n);divcent(cent);
	for(int i=1;i<=n;i++) if(a[i]==X) res1[i]++,res2[i]++;
//	for(int i=1;i<=n;i++) printf("%d %d\n",res1[i],res2[i]);
	ll res=0;
	for(int i=1;i<=n;i++){
		int rst1=n-res1[i],rst2=n-res2[i];
		res+=2ll*rst1*res1[i];res+=1ll*rst1*res2[i];
		res+=1ll*rst2*res1[i];res+=2ll*rst2*res2[i];
	} printf("%lld\n",1ll*n*n*n-(res>>1));
	return 0;
}