14940： [Ynoi2016]这是我自己的发明

#include<cstdio>
#include<algorithm>
#include<vector>
#include<map>
using namespace std;const int N=1e5+10;const int BC=230;typedef long long ll;
namespace INPUT_SPACE{
    const int BS=(1<<24)+5;char Buffer[BS],*HD,*TL;
    char gc() { if(HD==TL) TL=(HD=Buffer)+fread(Buffer,1,BS,stdin);return (HD==TL)?EOF:*HD++; }
    inline void read(int& t)
    {
        int x,ch,sgn=1;while(((ch=gc())<'0'||ch>'9')&&ch!='-');
        if(ch=='-') sgn=-1,x=0;else x=ch^'0';
        while((ch=gc())>='0'&&ch<='9') x=(x<<1)+(x<<3)+(ch^'0');
        t=x*sgn;
    }
}using INPUT_SPACE::read;
const int HB=1<<25;const int TB=(1<<25)-1;
template <typename T>
inline void gm(T* &bas,int siz,T* & op){op=bas;bas=bas+siz+1;}
int dfn[N];int nfd[N];int siz[N];int df;int dg[N];int col[N];int nrt;
struct qry{int u;int v;ll ans;}qr[N*5];int n;int m;int ccnt[N];int bel[N];
inline int isson(int u,int rt){return (nfd[u]<=nfd[rt])&&(nfd[rt]<nfd[u]+siz[u]);}
namespace lsh
{
    map <int,int> mp;
    inline void pre()
    {
        for(int i=1;i<=n;i++)mp[col[i]]=1;map <int,int> :: iterator it,it1;
        for(it=mp.begin(),it1=it,it1++;it1!=mp.end();++it,++it1)it1->second+=it->second;
        for(int i=1;i<=n;i++)col[i]=mp[col[i]];for(int i=1;i<=n;i++)ccnt[col[i]]++;
    }
}
namespace brusolve1
{
    int* v[N];int Eg_bas[N<<2];int* A_t;int csiz[N];int afn[N];int af;
    inline void ih(){A_t=Eg_bas;for(int i=1;i<=n;i++)gm(A_t,dg[i]-(i!=1),v[i]),dg[i]=0;}
    inline void ins(int u,int V){v[u][++dg[u]]=V;}
    inline void pb(int u){afn[++af]=u;}
    inline void calc(int cc)
    {
        for(int i=1;i<=n;i++)csiz[i]=(col[i]==cc);
        for(int i=1;i<=n;i++)
        {int u=afn[i];for(int j=1;j<=dg[u];j++)csiz[u]+=csiz[v[u][j]];}
        for(int i=1,tot=csiz[1];i<=m;i++)
        {
            qry& p=qr[i];
            int c1=(p.u&HB)?tot-csiz[p.u&TB]:csiz[p.u&TB];
            int c2=(p.v&HB)?tot-csiz[p.v&TB]:csiz[p.v&TB];p.ans+=(ll)c1*c2;
        }
    }inline void solve(){for(int i=1;i<=n;i++)if(ccnt[i]>BC)calc(i);}
}
namespace brusolve2
{
    const int B=270;
    struct sqry{int l;int r;int id;};vector <sqry> ve[N];vector <sqry> ed; 
    int* vc[N];int V_bas[N<<2];int* A_t;int a1[N];int sub[N];int bi[N];
    inline int cqry(int l,int r)
    {
        if(bi[l]==bi[r]){int ret=0;for(int i=l;i<=r;i++)ret+=a1[i];return ret;}
        int ret=0;for(int i=l;bi[i]==bi[l];i++)ret+=a1[i];
        for(int i=r;bi[i]==bi[r];i--)ret+=a1[i];
        for(int i=bi[l]+1;i<bi[r];i++)ret+=sub[i];return ret;
    }
    inline void pre()
    {
        A_t=V_bas;for(int i=1;i<=n;i++)
            if(ccnt[i]<=BC){gm(A_t,ccnt[i],vc[i]),ccnt[i]=0;}
        for(int i=1;i<=n;i++)
            {int nc=col[i];if(ccnt[nc]<=BC)vc[nc][++ccnt[nc]]=nfd[i];}
        for(int i=1;i<=n;i++)bi[i]=(i-1)/B+1;
        for(int i=1;i<=m;i++)
        {
            qry& t=qr[i];int tp=(((t.u&HB)>>25)<<1)|((t.v&HB)>>25);
            int u=t.u&TB;int v=t.v&TB;sqry nw=(sqry){nfd[v],nfd[v]+siz[v]-1,0};
            vector <sqry>& q1=ve[nfd[u]-1];vector <sqry>& q2=ve[nfd[u]+siz[u]-1];
            switch(tp)
            {
                case 0:{nw.id=i<<2|1;q1.push_back(nw);nw.id=i<<2|0;q2.push_back(nw);break;}
                case 1:{nw.id=i<<2|3;q1.push_back(nw);nw.id=i<<2|2;q2.push_back(nw);break;}
                case 2:
                {
                    nw.id=i<<2|0;q1.push_back(nw);nw.id=i<<2|1;q2.push_back(nw);
                    ed.push_back((sqry){nfd[v],nfd[v]+siz[v]-1,i<<2|0});break;
                }
                case 3:
                {
                    nw.id=i<<2|2;q1.push_back(nw);nw.id=i<<2|3;q2.push_back(nw);
                    ed.push_back((sqry){nfd[v],nfd[v]+siz[v]-1,i<<2|2});break;
                }
            }
        }
    }
    inline void solve()
    {
        pre();ll tot=0;
        for(int i=1;i<=n;i++)
        {
            int nc=col[dfn[i]];
            if(ccnt[nc]<=BC)for(int j=1;j<=ccnt[nc];j++)
                    a1[vc[nc][j]]++,sub[bi[vc[nc][j]]]++,tot++;
            for(vector <sqry>:: iterator it=ve[i].begin();it!=ve[i].end();++it)
            {
                int ans=cqry(it->l,it->r);ll& tar=qr[it->id>>2].ans;
                switch(it->id&3)
                {
                    case 0:{tar+=ans;break;}case 1:{tar-=ans;break;}
                    case 2:{tar+=tot-ans;break;}case 3:{tar-=tot-ans;break;}
                }
            }
        }for(int i=1;i<=n;i++)a1[i]+=a1[i-1];
        for(vector <sqry>:: iterator it=ed.begin();it!=ed.end();++it)
        {
            int ans=a1[it->r]-a1[it->l-1];ll& tar=qr[it->id>>2].ans;
            if((it->id&3)==0)tar+=ans;else tar+=tot-ans;
        }
    }
}
namespace oldtree
{
    int v[N<<1];int x[N<<1];int ct;int al[N];int fa[N][22];int dep[N];
    inline void add(int u,int V){v[++ct]=V;x[ct]=al[u];al[u]=ct;dg[u]++;dg[V]++;}
    inline int dfs(int u,int f)
    {
        for(int i=0;fa[u][i];i++)fa[u][i+1]=fa[fa[u][i]][i];
        dfn[++df]=u;nfd[u]=df;
        for(int i=al[u];i;i=x[i])
            if(v[i]!=f)fa[v[i]][0]=u,dep[v[i]]=dep[u]+1,brusolve1::ins(u,v[i]),siz[u]+=dfs(v[i],u);
        brusolve1::pb(u);return ++siz[u];
    }
    inline int cbt(int u,int v)
    {
        int del=dep[v]-dep[u];del--;
        for(int i=0;del;del>>=1,i++)if(del&1)v=fa[v][i];return v;
    }
    inline void build(){brusolve1::ih();dfs(1,0);}
}
int main()
{
    read(n);read(m);
    for(int i=1;i<=n;i++)read(col[i]);lsh::pre();
    for(int i=1,u,v;i<n;i++)read(u),read(v),oldtree::add(u,v),oldtree::add(v,u);nrt=1;
    oldtree::build();
    for(int i=1,tp,u,v;i<=m;i++)
    {
        read(tp);if(tp==1){read(nrt);m--,i--;continue;}
        read(u);read(v);
        if(u==nrt)u=1;else if(isson(u,nrt))u=oldtree::cbt(u,nrt)|(1<<25);
        if(v==nrt)v=1;else if(isson(v,nrt))v=oldtree::cbt(v,nrt)|(1<<25);
        qr[i]=(qry){u,v,0};
    }brusolve1::solve();brusolve2::solve();
    for(int i=1;i<=m;i++)printf("%lld\n",qr[i].ans);return 0;
}