题解 P4689 【[Ynoi2016]这是我自己的发明】

· · 题解

P5268 [SNOI2017]一个简单的询问的树上加强版，难度得到了一定的提高

题解：

首先不考虑换根操作的话，询问可以用dfs序转为两个区间，分别对应两棵子树，于是问题就转化为了P5268

然后考虑换根操作，延续不换根的做法，我们现在也要搞出x,y在的dfs序上的对应区间

设rt为当前根，很容易发现只有当rt在x(或y)的子树内的时候才会有异于不换根做法的情况出现

画出这种情况：

稍加观察就能发现以rt为根时，x的子树是整棵树扣掉1的子树，而1是rt的祖先中最接近x的一个

然后类比，所有的这种"1"都可以用树上k级祖先O(1)-O(log n)解决，于是x的对应区间就是整棵树扣掉"1"的子树区间

搞出每个询问对应的实际区间后就可以仿照P5268的做法拆询问搞莫队啦qwq

代码：

#pragma GCC optimize(3,"Ofast","inline")
#pragma GCC target("avx,avx2")
#include <bits/stdc++.h>
using namespace std;
#ifdef ONLINE_JUDGE
#define getchar gc
inline char gc(){
    static char buf[1<<16],*p1=buf,*p2=buf;
    if(p1==p2){
        p2=(p1=buf)+fread(buf,1,1<<16,stdin);
        if(p2==p1) return EOF;
    }
    return *p1++;
}
#endif
template<class t> inline t read(t &x){
    char c=getchar();bool f=0;x=0;
    while(!isdigit(c)) f|=c=='-',c=getchar();
    while(isdigit(c)) x=(x<<1)+(x<<3)+(c^48),c=getchar();
    if(f) x=-x;return x;
}
template<class t> inline void write(t x){
    if(x<0) putchar('-'),write(-x);
    else{if(x>9) write(x/10);putchar('0'+x%10);}
}

#define pii pair<int,int>
#define l first
#define r second

const int N=5e5+5,M=1e6+5; //理性开大，没啥坏处（雾）
int dfn[N],cnt,sz[N],d[N],f[N][18],qn,ans[M],t,rt,bl[N],n,m,num[N],un,a[N],col[N],cntl[N],cntr[N];
vector<int> g[N];

struct que{
    int l,r,i;
    inline bool operator < (const que &nt) const {
        return bl[l]^bl[nt.l]?bl[l]<bl[nt.l]:bl[l]&1?r<nt.r:r>nt.r;
    }
}q[M<<2];

void dfs(int x,int fa){ //预处理
    dfn[x]=++cnt;sz[x]=1;d[x]=d[fa]+1;
    f[x][0]=fa;
    for(int i=1;f[f[x][i-1]][i-1];i++) f[x][i]=f[f[x][i-1]][i-1];
    for(int y:g[x]) if(y^fa) dfs(y,x),sz[x]+=sz[y];
}

void add(int l,int r,int ll,int rr,int i){ //拆询问，类比P5268
    q[++qn]=(que){r,rr,i};
    q[++qn]=(que){l-1,rr,-i}; 
    q[++qn]=(que){ll-1,r,-i};
    q[++qn]=(que){l-1,ll-1,i};
}

int getk(int x,int k){ //logn找k祖先
    for(int i=17;~i;i--) if(k>=(1<<i)) k-=1<<i,x=f[x][i];
    return x;
}

vector<pii> calc(int x){
    vector<pii> res;
    if(x==rt) res.push_back(pii(1,n)); //是根就是[1,n]
    else if(dfn[x]<=dfn[rt]&&dfn[rt]<=dfn[x]+sz[x]-1){ //rt在x子树内就找出要被扣掉的那棵树的根节点son
        int son=getk(rt,d[rt]-d[x]-1);
        if(1<=dfn[son]-1) res.push_back(pii(1,dfn[son]-1));
        if(dfn[son]+sz[son]<=n) res.push_back(pii(dfn[son]+sz[son],n));
    }
    else res.push_back(pii(dfn[x],dfn[x]+sz[x]-1)); //否则就直接返回x子树的dfs序区间
    return res;
}

signed main(){
    srand(time(NULL));
    read(n);read(m);rt=1;
    for(int i=1;i<=n;i++) num[i]=read(a[i]);
    sort(num+1,num+1+n);
    un=unique(num+1,num+1+n)-num-1;
    for(int i=1,x,y;i<n;i++){
        read(x);read(y);
        g[x].push_back(y);
        g[y].push_back(x);
    }
    dfs(rand()%n+1,0); //小优化，随机为根让毒瘤出题人很难来卡你
    for(int i=1;i<=n;i++) col[dfn[i]]=lower_bound(num+1,num+1+un,a[i])-num; //将原数列点权转为对应dfs序列的点权，并顺便离散化
    while(m--){
        int op;
        read(op);
        if(op==1) read(rt);
        else{
            int x,y;
            vector<pii> px,py; //分别对应两个询问点的对应ds区间
            t++;
            read(x);read(y);
            px=calc(x);
            py=calc(y);
            for(auto x:px) for(auto y:py) add(x.l,x.r,y.l,y.r,t); //分别枚举，进行拆分询问
        }
    }
    //下面就与P5268相同了
    int len=n/sqrt(t);
    for(int i=1;i<=n;i++) bl[i]=(i-1)/len+1;
    for(int i=1;i<=qn;i++) if(q[i].l>q[i].r) swap(q[i].l,q[i].r);
    sort(q+1,q+1+qn);
    for(int i=1,l=1,r=1,cur=0;i<=qn;i++){
        while(l<q[i].l) cur+=cntr[col[++l]],cntl[col[l]]++;
        while(l>q[i].l) cur-=cntr[col[l]],cntl[col[l--]]--;
        while(r<q[i].r) cur+=cntl[col[++r]],cntr[col[r]]++;
        while(r>q[i].r) cur-=cntl[col[r]],cntr[col[r--]]--;
        ans[abs(q[i].i)]+=cur*(q[i].i/abs(q[i].i));
    }
    for(int i=1;i<=t;i++) write(ans[i]),puts("");
}