Codeforces 986E - Prince's Problem（树上前缀和）

题面传送门

题意：

有一棵 \(n\) 个节点的树，点上有点权 \(a_i\)，\(q\) 组询问，每次询问给出 \(u,v,w\)，要求：

• \(\prod\limits_{x\in P(u,v)}\operatorname{gcd}(a_x,w)\)

其中 \(P(u,v)\) 为 \(u\) 到 \(v\) 的简单路径上经过的点的集合。

\(n,q \leq 10^5,a_i \leq 10^7\)

显然本题需要分解质因数，不妨设 \(w\) 分解质因数得到 \(p_1^{f_1}\times p_2^{f_2}\times\dots\times p_k^{f_k}\)。

我们对 \(w\) 的每个质因子 \(p_i\) 计算对应的贡献，对于链上的每个数 \(a_x\)，假设它质因数 \(p_i\) 的次数为 \(g_i\)，那么它的贡献就是 \(p_i^{\min(f_i,g_i)}\)。

于是我们可以暴力枚举 \(p_i\) 的次数 \(g_i\)，假设 \(u\to v\) 的路径上有 \(cnt\) 个 \(a_x\) 中 \(p_i\) 的次数恰好为 \(g_i\)，这部分的答案应为 \((p_i^{\min(f_i,g_i)})^{cnt}\)。

接下来就是怎样求这个 \(cnt\)，对 \(a_i\) 分解质因数得到 \(a_i=q_1^{h_1}\times q_2^{h_2}\times\dots\times q_l^{h_l}\)，将每个 \((q_i,h_i)\) 看作一种“颜色”，每个点上最多有 \(7\) 种颜色。那么题目就转化为求 \(u\) 到 \(v\) 路径上有某种“颜色”的点的个数。

考虑用树剖把原区间拆成 \(\log n\) 个小区间，然后用主席树维护区间内值为 \(x\) 的数的个数就可以求出 \(cnt\) 的值。

算算复杂度，暴力枚举 \(g_i\) 的常数最大可以达到为 \(\log_2(10^7)+\log_3(10^7)+\log_5(10^7)+\log_7(10^7)+\log_{11}(10^7)+\log_{13}(10^7)\) ，大约 \(60\)，树剖+主席树各有一个 \(\log\)，所以总复杂度是 \(60q\log^2n\)，太逊了。

发现本题并没有强制在线，所以可以离线地去完成路径上数颜色这一工作。这玩意儿可以利用树上差分拆为 \(u,v,\operatorname{lca}(u,v),fa_{\operatorname{lca}(u,v)}\) 四个部分，最终的 \(cnt=cnt_u+cnt_v-cnt_{\operatorname{lca}(u,v)}-cnt_{fa_{\operatorname{lca}(u,v)}}\)。

枚举每一种颜色 \(c\)，然后依次处理查询颜色为 \(c\) 的询问。假设某个点 \(x\) 拥有颜色 \(c\)，那么它会对 \(x\) 的子树内的询问产生 \(1\) 的贡献，利用 DFS 序将子树变为 DFS 序上的一个区间，于是本题就变为区间加、单点查询的问题，一个线段树就可以搞定了。

复杂度可以少掉一个 \(\log\)， \(60q\log n\) 大约为 \(1.5\times 10^8\)。

想到这里，我就没有再想下去。当我美滋滋地交上去，拥有十足的信心 A 掉这道题的时候……

TLE on test 9？

虽然理论复杂度是 ok 的，但实测常数很大。首先你要将每个找颜色的询问拆成四个询问，就已经有了 \(4\) 的常数，再加上快速幂、取模的常数，将近有 \(10^9\)。对这种情况，哪怕再神的评测机也没法儿给你跑过啊……

顺便附上 T 掉了的代码：

#include <bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define fz(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
#define ffe(it,v) for(__typeof(v.begin()) it=v.begin();it!=v.end();it++)
#define fill0(a) memset(a,0,sizeof(a))
#define fill1(a) memset(a,-1,sizeof(a))
#define fillbig(a) memset(a,63,sizeof(a))
#define pb push_back
#define ppb pop_back
#define mp make_pair
typedef pair<int,int> pii;
typedef long long ll;
const int MOD=1e9+7;
const int MAXN=1e5+5;
const int MAXC=MAXN*10;
const int MAXQ=MAXN*150;
int n,q,qn,a[MAXN],tnum=0,num=0;
vector<pii> fac[MAXN];
pii key[MAXC],hs[MAXC];
map<pii,bool> vis;
int hd[MAXN],nxt[MAXN<<1],to[MAXN<<1],ec=0;
void adde(int u,int v){
	to[++ec]=v;nxt[ec]=hd[u];hd[u]=ec;
}
int fa[MAXN][22],dfn[MAXN],ed[MAXN],tim=0,dep[MAXN];
void dfs(int x,int f){
	fa[x][0]=f;dfn[x]=++tim;
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==f) continue;
		dep[y]=dep[x]+1;dfs(y,x);
	} ed[x]=tim;
}
int getlca(int x,int y){
	if(dep[x]<dep[y]) swap(x,y);
	for(int i=20;~i;i--) if(dep[x]-(1<<i)>=dep[y]) x=fa[x][i];
	if(x==y) return x;
	for(int i=20;~i;i--) if(fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i];
	return fa[x][0];
}
struct Query{
	int f,id,x,v;//ans[id]*=v^{num}
	Query(){/*var nothing*/}
	Query(int _f,int _id,int _x,int _v){
		f=_f;id=_id;x=_x;v=_v;
	}
	friend bool operator <(Query lhs,Query rhs){
		return lhs.f<rhs.f;
	}
} qu[MAXQ];
vector<int> add[MAXC];
int qpow(int x,int e){
	int ret=1;
	while(e){
		if(e&1) ret=1ll*ret*x%MOD;
		x=1ll*x*x%MOD;e>>=1;
	} return ret;
}
int ans[MAXN];
struct node{
	int l,r,val,lz;
} s[MAXN<<2];
void build(int k,int l,int r){
	s[k].l=l;s[k].r=r;if(l==r) return;
	int mid=(l+r)>>1;build(k<<1,l,mid);build(k<<1|1,mid+1,r);
}
void pushdown(int k){
	if(s[k].lz){
		s[k<<1].val+=s[k].lz*(s[k<<1].r-s[k<<1].l+1);
		s[k<<1].lz+=s[k].lz;
		s[k<<1|1].val+=s[k].lz*(s[k<<1|1].r-s[k<<1|1].l+1);
		s[k<<1|1].lz+=s[k].lz;
		s[k].lz=0;
	}
}
void modify(int k,int l,int r,int x){
	if(l<=s[k].l&&s[k].r<=r){
		s[k].val+=(s[k].r-s[k].l+1)*x;s[k].lz+=x;return;
	} pushdown(k);int mid=(s[k].l+s[k].r)>>1;
	if(r<=mid) modify(k<<1,l,r,x);
	else if(l>mid) modify(k<<1|1,l,r,x);
	else modify(k<<1,l,mid,x),modify(k<<1|1,mid+1,r,x);
	s[k].val=s[k<<1].val+s[k<<1|1].val;
}
int query(int k,int l,int r){
	if(l<=s[k].l&&s[k].r<=r) return s[k].val;
	pushdown(k);int mid=(s[k].l+s[k].r)>>1;
	if(r<=mid) return query(k<<1,l,r);
	else if(l>mid) return query(k<<1|1,l,r);
	else return query(k<<1,l,mid)+query(k<<1|1,mid+1,r);
}
int main(){
	scanf("%d",&n);
	for(int i=1;i<n;i++){
		int u,v;scanf("%d%d",&u,&v);adde(u,v);adde(v,u);
	}
	for(int i=1;i<=n;i++) scanf("%d",&a[i]);
	for(int i=1;i<=n;i++){
		int tmp=a[i];
		for(int j=2;j*j<=tmp;j++){
			if(tmp%j==0){
				int cnt=0;
				while(tmp%j==0) tmp/=j,cnt++;
				fac[i].pb(mp(j,cnt));
//				printf("{%d,%d} ",j,cnt);
				key[++tnum]=mp(j,cnt);
			}
		}
		if(tmp>1){
			fac[i].pb(mp(tmp,1));
//			printf("{%d,%d} ",tmp,1);
			key[++tnum]=mp(tmp,1);
		}
//		printf("\n");
	}
	sort(key+1,key+tnum+1);
	for(int i=1;i<=tnum;i++) if(key[i]!=key[i-1])
		hs[++num]=key[i],vis[key[i]]=1;
	for(int i=1;i<=n;i++){
		for(int j=0;j<fac[i].size();j++){
			int id=lower_bound(hs+1,hs+num+1,fac[i][j])-hs;
			add[id].push_back(i);//printf("%d %d\n",id,i);
		}
	}
	dfs(1,0);
	for(int i=1;i<=20;i++) for(int j=1;j<=n;j++)
		fa[j][i]=fa[fa[j][i-1]][i-1];
	scanf("%d",&q);
	for(int i=1;i<=q;i++){
		int u,v,w;scanf("%d%d%d",&u,&v,&w);
		int tmp=w;vector<pii> ff;ans[i]=1;
		for(int j=2;j*j<=tmp;j++){
			if(tmp%j==0){
				int cnt=0;
				while(tmp%j==0) tmp/=j,cnt++;
				ff.pb(mp(j,cnt));
			}
		}
		if(tmp>1) ff.pb(mp(tmp,1));
		for(int j=0;j<ff.size();j++){
			int val=1,pw=qpow(ff[j].fi,ff[j].se);
			for(int k=1;val<=1e7;k++){
				val*=ff[j].fi;
				if(!vis[mp(ff[j].fi,k)]) continue;
				int f=lower_bound(hs+1,hs+num+1,mp(ff[j].fi,k))-hs;
				qu[++qn]=Query(f,i,u,min(val,pw));
				qu[++qn]=Query(f,i,v,min(val,pw));
				int lca=getlca(u,v);
				qu[++qn]=Query(f,i,lca,qpow(min(val,pw),MOD-2));
				if(lca!=1) qu[++qn]=Query(f,i,fa[lca][0],qpow(min(val,pw),MOD-2));
			}
		}
	}
	build(1,1,n);sort(qu+1,qu+qn+1);int j=1;
	for(int i=1;i<=num;i++){
		ffe(it,add[i]) modify(1,dfn[*it],ed[*it],1);
		while(j<=qn&&qu[j].f==i){
			int cnt=query(1,dfn[qu[j].x],dfn[qu[j].x]);
//			printf("%d %d\n",j,cnt);
			ans[qu[j].id]=1ll*ans[qu[j].id]*qpow(qu[j].v,cnt)%MOD;j++;
		}
		ffe(it,add[i]) modify(1,dfn[*it],ed[*it],-1);
	}
	for(int i=1;i<=q;i++) printf("%d\n",ans[i]);
	return 0;
}

wtcl 啊，下面开始讲正解，其实稍微变通一下就行了。

正解还是基于之前那个“数颜色”的做法，不过我们想想，之前我们离线是枚举颜色，然后维护树上每个点询问的答案。如果是枚举树上的点，然后用个什么东西维护每种颜色的数量会不会好一些啊？

树上前缀和啊！

按照套路对树进行一遍 DFS，将询问标记在点上。我们维护一个 \(sum_x\) 表示根节点到当前节点的路径上颜色为 \(x\) 的节点的个数。在 DFS 的过程中实时维护 \(sum\) 数组，回溯的时候就减掉当前节点的贡献。

前缀和可以 \(\mathcal O(1)\) 的时间内维护，这样又少掉一个 \(\log\)，可以通过此题。

#include <bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define fz(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
#define ffe(it,v) for(__typeof(v.begin()) it=v.begin();it!=v.end();it++)
#define fill0(a) memset(a,0,sizeof(a))
#define fill1(a) memset(a,-1,sizeof(a))
#define fillbig(a) memset(a,63,sizeof(a))
#define pb push_back
#define ppb pop_back
#define mp make_pair
typedef pair<int,int> pii;
typedef long long ll;
const int MOD=1e9+7;
const int MAXN=1e5+5;
const int MAXV=1e7+5;
int n,q,qn,a[MAXN];
int hd[MAXN],nxt[MAXN<<1],to[MAXN<<1],ec=0;
void adde(int u,int v){
	to[++ec]=v;nxt[ec]=hd[u];hd[u]=ec;
}
int fa[MAXN][22],dep[MAXN];
void dfs(int x,int f){
	fa[x][0]=f;
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==f) continue;
		dep[y]=dep[x]+1;dfs(y,x);
	}
}
int getlca(int x,int y){
	if(dep[x]<dep[y]) swap(x,y);
	for(int i=20;~i;i--) if(dep[x]-(1<<i)>=dep[y]) x=fa[x][i];
	if(x==y) return x;
	for(int i=20;~i;i--) if(fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i];
	return fa[x][0];
}
vector<pii> fac[MAXN];
struct Query{
	int w,x,id,mul;
} qu[MAXN<<2];
int qpow(int x,int e){
	int ret=1;
	while(e){
		if(e&1) ret=1ll*ret*x%MOD;
		x=1ll*x*x%MOD;e>>=1;
	} return ret;
}
int ans[MAXN];
int sum[MAXV/5][24];
vector<int> qv[MAXN];
int pid[MAXV],pr[MAXV/5],pcnt=0;
bitset<MAXV> vis;
void prework(){
	for(int i=2;i<MAXV;i++){
		if(!vis[i]){pr[++pcnt]=i;pid[i]=pcnt;}
		for(int j=1;j<=pcnt&&pr[j]*i<MAXV;j++){
			vis[pr[j]*i]=1;if(i%pr[j]==0) break;
		}
	}
}
void calc(int x){
	ffe(it,fac[x]) sum[pid[it->fi]][it->se]++;
	ffe(it,qv[x]){
		int id=*it,tmp=qu[id].w;
		vector<pii> ff;
		for(int i=2;i*i<=tmp;i++){
			if(tmp%i==0){
				int cnt=0;
				while(tmp%i==0) tmp/=i,cnt++;
				ff.pb(mp(i,cnt));
			}
		}
		if(tmp>1) ff.pb(mp(tmp,1));
		int tot_mul=1;
		for(int i=0;i<ff.size();i++){
			int pw=1,val=ff[i].fi;
			for(int j=1;j<=ff[i].se;j++) pw*=ff[i].fi;
			for(int j=1;val<MAXV;j++,val*=ff[i].fi){
				int cnt=sum[pid[ff[i].fi]][j];
				int mul=qpow(min(val,pw),cnt);
				tot_mul=1ll*tot_mul*mul%MOD;
			}
		}
		if(qu[id].mul==1) ans[qu[id].id]=1ll*ans[qu[id].id]*tot_mul%MOD;
		else ans[qu[id].id]=1ll*ans[qu[id].id]*qpow(tot_mul,MOD-2)%MOD;
	}
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==fa[x][0]) continue;calc(y);
	}
	ffe(it,fac[x]) sum[pid[it->fi]][it->se]--;
}
int main(){
	prework();
	scanf("%d",&n);
	for(int i=1;i<n;i++){
		int u,v;scanf("%d%d",&u,&v);adde(u,v);adde(v,u);
	}
	for(int i=1;i<=n;i++) scanf("%d",&a[i]);
	for(int i=1;i<=n;i++){
		int tmp=a[i];
		for(int j=2;j*j<=tmp;j++){
			if(tmp%j==0){
				int cnt=0;
				while(tmp%j==0) tmp/=j,cnt++;
				fac[i].pb(mp(j,cnt));
			}
		}
		if(tmp>1) fac[i].pb(mp(tmp,1));
	}
	dfs(1,0);
	for(int i=1;i<=20;i++) for(int j=1;j<=n;j++)
		fa[j][i]=fa[fa[j][i-1]][i-1];
	scanf("%d",&q);
	for(int i=1;i<=q;i++){
		int u,v,w;scanf("%d%d%d",&u,&v,&w);ans[i]=1;
		int lca=getlca(u,v);
		qu[++qn]={w,u,i,1};qu[++qn]={w,v,i,1};
		qu[++qn]={w,lca,i,-1};if(lca!=1) qu[++qn]={w,fa[lca][0],i,-1};
	}
	for(int i=1;i<=qn;i++) qv[qu[i].x].pb(i);calc(1);
	for(int i=1;i<=q;i++) printf("%d\n",ans[i]);
	return 0;
}