bzoj4940 树刨+莫队(好题)

题意：
一颗树n个点有点权，m次操作。
操作两种：

1. 换根
2. 查询 x点的子树选每一个点，y点的子树选每一个点，如果两点点权相同ans++，输出ans
   n1e5,m5e5

换根是幌子，先按照1来树刨，这个跟bzoj3083结论一样，就是分三种情况：

1. 当前根节点与被查询的点一样，那么这个点的子树区间直接是1-n
2. 当前根节点与被查询的点x的LCA与x不一样，即lca(root,x)!=x，查询区间是（p[x]+sz[x]-1）
3. 第三种就是根节点与x的LCA一样，及lca(root,x)==x，查询区间是（1-p[y]-1)+(p[y]+sz[y],n)，y是x的儿子，且自身或包含root
   将查询子树换为查询区间之后，根bzoj5016结论一样，把这四个区间最多转换为9个区间。结论搜一搜这个题号，有讲的详细的：

然后就出来了：

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define forn(i,n) for(int i=0;i<n;i++)
#define for1(i,n) for(int i=1;i<=n;i++)
#define IO ios::sync_with_stdio(false);cin.tie(0)
const int maxn = 1e5+5;
const int maxm = 5e5+5;
 
int a[maxn],c[maxn],pos[maxn],sz[maxn],par[maxn],deep[maxn],son[maxn],top[maxn],p[maxn],cnt1[maxn],cnt2[maxn],fp[maxn],id,root;
vector<int>e[maxn],b;
ll ans[maxm];
struct Q{
    int l,r,flag,id;
}q[maxm*9];
bool cmp(Q a,Q b){
    if(pos[a.l]==pos[b.l]) return a.r<b.r;
    return pos[a.l]<pos[b.l];
}
void dfs(int u,int pre,int d){
    deep[u] = d,par[u] = pre,sz[u] = 1;
    forn(i,e[u].size()){
        int v = e[u][i];
        if(v==pre) continue;
        dfs(v,u,d+1);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u] = v;
    }
}
void getpos(int u,int gg){
    top[u] = gg,p[u] = ++id,fp[p[u]] = u;
    if(son[u]) getpos(son[u],gg);
    forn(i,e[u].size()){
        int v = e[u][i];
        if(v==par[u]||v==son[u]) continue;
        getpos(v,v);
    }
}
int LCA(int x,int y){
    int fx = top[x],fy = top[y];
    while(fx!=fy){
        if(deep[fx]<deep[fy]) swap(fx,fy),swap(x,y);  
        x = par[fx],fx = top[x];
    }
    if(deep[x]>deep[y]) swap(x,y);
    return x;
}
 
int main(){
    IO;
    int n,m;cin>>n>>m;
    int zz = sqrt(n);
    for1(i,n) {
        cin>>a[i];
        pos[i] = i/zz;
        b.push_back(a[i]);
    }
    sort(b.begin(),b.end());
    b.erase(unique(b.begin(),b.end()),b.end());
    for1(i,n){
        int pp = lower_bound(b.begin(),b.end(),a[i])-b.begin();
        a[i] = pp;
    }
    forn(i,n-1){
        int x,y;cin>>x>>y;
        e[x].push_back(y);
        e[y].push_back(x);
    }
    root = 1;
    int cnt = 0,cntq = 0;
    dfs(1,0,0);
    getpos(1,1);
    for1(i,n) c[i] = a[fp[i]];
    //for1(i,n) cerr<<c[i]<<' ';
   // cerr<<'\n';
    //for1(i,n) cerr<<p[i]<<' ';
    //cerr<<'\n';
    forn(i,m){
        int op,x;cin>>op>>x;
        if(op==1) root = x;
        else{
            int y; cin>>y;
            vector<pair<int,int> >aa,bb;
            int lca = LCA(x,root);
            //cerr<<lca<<'\n';
            if(x==root) aa.push_back({1,n});
            else if(x!=lca) aa.push_back({p[x],p[x]+sz[x]-1});
            else{
                int z;
                forn(j,e[x].size()){
                    int v = e[x][j];
                    if(p[v]>=p[x]&&p[v]<=p[root]&&p[v]+sz[v]-1>=p[root]){
                        z = v;
                        break;  
                    }
                }
                if(p[z]!=1)aa.push_back({1,p[z]-1});
                if(p[z]+sz[z]-1!=n) aa.push_back({p[z]+sz[z],n});
            }
            lca = LCA(y,root);
            if(y==root) bb.push_back({1,n});
            else if(y!=lca) bb.push_back({p[y],p[y]+sz[y]-1});
            else{
                int z;
                forn(j,e[y].size()){
                    int v = e[y][j];
                    if(p[v]>=p[y]&&p[v]<=p[root]&&p[v]+sz[v]-1>=p[root]){
                        z = v;
                        break;  
                    }
                }
                if(p[z]!=1)bb.push_back({1,p[z]-1});
                if(p[z]+sz[z]-1!=n) bb.push_back({p[z]+sz[z],n});
            }
            forn(xxx,aa.size()){
                forn(yyy,bb.size()){
                    pair<int,int> xx = aa[xxx],yy=bb[yyy];
                    int l1 = xx.first,r1 = xx.second,l2 = yy.first,r2 = yy.second;
                    //cerr<<"@!#!@#!@#!@#!@#       "<<l1<<' '<<r1<<' '<<l2<<' '<<r2<<' '<<cntq<<'\n';
                    q[cnt++] = {min(r1,r2),max(r1,r2),1,cntq};
                    if(l1>1) q[cnt++] = {min(l1-1,r2),max(l1-1,r2),-1,cntq};
                    if(l2>1) q[cnt++] = {min(l2-1,r1),max(l2-1,r1),-1,cntq};
                    if(l1>1&&l2>1) q[cnt++] = {min(l1-1,l2-1),max(l1-1,l2-1),1,cntq};
                }
            }
            cntq++;
        }
    }
    sort(q,q+cnt,cmp);
    ll res = 0;
    int l = 0,r = 0;
  //  for1(i,n) cerr<<a[i]<<' ';
  //  cerr<<'\n';
  //  for1(i,n) cerr<<c[i]<<' ';
   // cerr<<'\n';
    //cerr<<"!@#"<<'\n';
    forn(i,cnt){
        //cerr<<q[i].l<<' '<<q[i].r<<'\n';
        while(r<q[i].r) r++,res+=cnt1[c[r]],cnt2[c[r]]++;
        while(r>q[i].r) cnt2[c[r]]--,res-=cnt1[c[r]],r--;
        //cerr<<l<<' '<<r<<' '<<c[r]<<' '<<cnt2[c[r]]<<' '<<res<<'\n';
        while(l<q[i].l) l++,res+=cnt2[c[l]],cnt1[c[l]]++;
       // cerr<<res<<'\n';
        while(l>q[i].l) cnt1[c[l]]--,res-=cnt2[c[l]],l--;
        //cerr<<q[i].l<<' '<<q[i].r<<' '<<q[i].id<<' '<<q[i].flag<<' '<<res<<'\n';
        ans[q[i].id] += q[i].flag*res;
    }
   // cout<<'\n';
    forn(i,cntq) cout<<ans[i]<<'\n';
    return 0;
}
/*
12 14
1 1 2 4 1 5 2 1 1 1 1 1
1 4
4 9 
4 7
4 5
5 3
1 11
10 12
6 8
6 10
11 2
11 6
 
2 11 4
2 4 1
2 5 4
1 4
2 9 5
2 5 4
2 4 4
2 4 1
1 6
2 6 6
2 10 8
2 1 11
2 2 11
2 6 11
*/