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;
}