树上莫队 洛谷SP10707 COT2 - Count on a tree II

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=200200;
int dep[N],fa[N][30],color[N],bin[N],blg[N],ord[N],in[N],out[N],ans[N],h[N<<1],cnt[N],vis[N],tot,id,now;
struct event
{
    int l,r,pos,lca;
}q[N];

struct edge
{
    int v,next;
}e[N<<1];

void addedge(int u,int v)
{
    e[++tot]=(edge){v,h[u]};
    h[u]=tot;
}
void build(int n)
{
    int sz=sqrt(n);
    for(int i=1;i<=n;i++)blg[i]=(i-1)/sz+1;
}
static bool cmp(event &a,event &b)
{
    return (blg[a.l]^blg[b.l])?(blg[a.l]<blg[b.l]):((blg[a.l]&1)?a.r<b.r:a.r>b.r);
}
inline ll read()
{
    ll x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=(x<<1)+(x<<3)+(ch^48);ch=getchar();}
    return x*f;
}
void dfs(int u,int p)
{
    
    in[u]=++id;
    ord[id]=u;
    dep[u]=dep[p]+1;
    fa[u][0]=p;
    for(int i=1;(1<<i)<dep[u];i++)fa[u][i]=fa[fa[u][i-1]][i-1];
    for(int i=h[u];i;i=e[i].next)
    {
        int v=e[i].v;
        if(v==p)continue;
        dfs(v,u);
    }
    out[u]=++id;
    ord[id]=u;
}
int lca(int u,int v)
{
    if(dep[u]>dep[v])swap(u,v);
    for(int i=20;i>=0;i--)if(dep[fa[v][i]]>=dep[u])v=fa[v][i];
    if(u==v)return u; 
    for(int i=20;i>=0;i--)if(fa[v][i]!=fa[u][i]){v=fa[v][i];u=fa[u][i];}
    return fa[u][0];
}
void change(int node)
{
    if(vis[node])
    {
        if((--cnt[color[node]])==0)now--;
    }
    else
    {
        if((++cnt[color[node]])==1)now++;
    }
    vis[node]^=1; 
}
int main()
{
    int n=read(),m=read();
    for(int i=1;i<=n;i++)color[i]=bin[i]=read();
    for(int i=1;i<n;i++)
    {
        int u=read(),v=read();
        addedge(u,v);
        addedge(v,u);
    }
    dfs(1,0);
    
    for(int i=1;i<=m;i++)
    {
        int l=read(),r=read();
        q[i].pos=i;
        int lca_=lca(l,r);
        if(in[l]>in[r])swap(l,r);
        if(l==lca_)
        {
            q[i].l=in[l];
            q[i].r=in[r];
            q[i].lca=0;    
        }
        else
        {
            q[i].l=out[l];
            q[i].r=in[r];
            q[i].lca=lca_;
        }
    
    }
    build(n*2);//分块要考虑欧拉序，不然会T；
    sort(q+1,q+1+m,cmp);
    sort(bin+1,bin+1+n);
    int bcnt=unique(bin+1,bin+1+n)-bin-1;
    for(int i=1;i<=n;i++)color[i]=lower_bound(bin+1,bin+1+bcnt,color[i])-bin;
    int l=1,r=0;
    for(int i=1;i<=m;i++)
    {
        int ql=q[i].l,qr=q[i].r,lca_=q[i].lca;
        while(l<ql)change(ord[l++]);
        while(l>ql)change(ord[--l]);
        while(r<qr)change(ord[++r]);
        while(r>qr)change(ord[r--]);
        if(lca_)change(lca_);
        ans[q[i].pos]=now;
        if(lca_)change(lca_);
    }
    for(int i=1;i<=m;i++)printf("%d\n",ans[i]);
    return 0;
}