具體思路看別的題解吧。這里只提兩個可能對常數和代碼長度有優化的處理方法。
I
把一個詢問拆成\(9\)個甚至\(16\)個莫隊詢問實在是有點珂怕。
發現詢問的一邊要么是一個區間,要么是\([1,n]\)挖去一個區間。
記\(pre_i=f_{[1,i],[1,n]}\),這個可以一遍預處理求出來。
簡單容斥一下:
\[f_{[l,r],[1,L)\cup(R,n]}=f_{[l,r],[1,n]}-f_{[l,r],[L,R]}=pre_r-pre_{l-1}-f_{[l,r],[L,R]} \]
\[f_{[1,l)\cup(r,n],[1,L)\cup(R,n]}=f_{[1,n],[1,n]}-f_{[l,r],[1,n]}-f_{[1,n],[L,R]}+f_{[l,r],[L,R]}=... \]
於是對於每個詢問,我們拆成\(4\)個莫隊詢問就夠了,因為只差求\(f_{[l,r],[L,R]}\)。
II
如何求某一個點往新根方向上的兒子?樹剖倍增什么的都太呆了。
預處理樹的dfn序,每個點開一個map按dfn序掛上它所有的兒子,查詢時拿着新根的dfn序直接lower_bound即可。
然后就可以2kb寫完這道不算太毒瘤的由乃題了。
#include<bits/stdc++.h>
#define LL long long
#define RG register
#define R RG int
#define G if(++ip==ie)if(fread(ip=buf,1,SZ,stdin))
using namespace std;
const int SZ=1<<19,N=1e5+9,M=5e5+9;
char buf[SZ],*ie=buf+SZ,*ip=ie-1;
inline int in(){
G;while(*ip<'-')G;
R x=*ip&15;G;
while(*ip>'-'){x*=10;x+=*ip&15;G;}
return x;
}
struct Qry{
int x,y,b,id;
inline bool operator<(const Qry&t)const{
return b!=t.b?b<t.b:(1&b)?y>t.y:y<t.y;
}
}t[4*M];
int dfn,a[N],b[N],c1[N],c2[N],l[N],r[N],he[N],ne[2*N],to[2*N];
LL pre[N],ans[M];
map<int,int>ch[N];
void dfs(R x,R f){//預處理dfn序
l[x]=++dfn;
for(R y,i=he[x];i;i=ne[i])
if((y=to[i])!=f)dfs(y,x),ch[x][r[y]]=y;
r[x]=dfn;
}
int main(){
R n=in(),m=in(),q=0,B=sqrt(n);
for(R i=1;i<=n;++i)
a[i]=b[i]=in();
for(R i=1,p=0;i<n;++i){
R x=in(),y=in();
ne[++p]=he[x];to[he[x]=p]=y;
ne[++p]=he[y];to[he[y]=p]=x;
}
dfs(1,0);
sort(a+1,a+n+1);
for(R i=1;i<=n;++i)b[i]=lower_bound(a+1,a+n+1,b[i])-a;
for(R i=1;i<=n;++i)++c1[a[l[i]]=b[i]];//預處理pre
for(R i=1;i<=n;++i)pre[i]=pre[i-1]+c1[a[i]];
for(R i=1,rt=1;i<=m;++i){
if(in()&1){rt=in();--i;--m;continue;}
R x=in(),y=in();
R tx=l[x]<=l[rt]&&r[rt]<=r[x];if(x==rt)x=1,tx=0;
R ty=l[y]<=l[rt]&&r[rt]<=r[y];if(y==rt)y=1,ty=0;
if(tx)x=ch[x].lower_bound(l[rt])->second;//map找兒子
if(ty)y=ch[y].lower_bound(l[rt])->second;
R lx=l[x]-1,ly=l[y]-1,rx=r[x],ry=r[y];
if(tx&&ty)ans[i]=pre[n];
if(tx)ans[i]+=(pre[ry]-pre[ly])*(tx==ty?-1:1);
if(ty)ans[i]+=(pre[rx]-pre[lx])*(tx==ty?-1:1);
t[++q]=(Qry){rx,ry,rx/B,tx==ty?i:-i};//四個莫隊詢問,帶上容斥系數
t[++q]=(Qry){rx,ly,rx/B,tx==ty?-i:i};
t[++q]=(Qry){lx,ry,lx/B,tx==ty?-i:i};
t[++q]=(Qry){lx,ly,lx/B,tx==ty?i:-i};
}
sort(t+1,t+q+1);
LL now=0;memset(c1+1,0,4*n);
for(R i=1,x=0,y=0;i<=q;++i){//莫隊
while(x<t[i].x)++c1[a[++x]],now+=c2[a[x ]];
while(x>t[i].x)--c1[a[ x]],now-=c2[a[x--]];
while(y<t[i].y)++c2[a[++y]],now+=c1[a[y ]];
while(y>t[i].y)--c2[a[ y]],now-=c1[a[y--]];
R j=t[i].id;j>0?ans[j]+=now:ans[-j]-=now;
}
for(R i=1;i<=m;++i)
printf("%lld\n",ans[i]);
return 0;
}