P2486 [SDOI2011]染色

#include <bits/stdc++.h>
#include <tr1/unordered_map>
using namespace std;
#define ll long long
#define re register
#define pb push_back
#define fi first
#define se second
const int N=1e6+10;
const int mod=998244353;
void read(int &a)
{
    a=0;int d=1;char ch;
    while(ch=getchar(),ch>‘9‘||ch<‘0‘)
        if(ch==‘-‘)
            d=-1;
    a=ch^48;
    while(ch=getchar(),ch>=‘0‘&&ch<=‘9‘)
        a=(a<<3)+(a<<1)+(ch^48);
    a*=d;
}
int val[N],top[N],id[N],rk[N],f[N],siz[N],son[N],dep[N],cnt,rt=1,last=-1;
vector <int> v[N];
void dfs1(int x)
{
    siz[x]=1,dep[x]=dep[f[x]]+1;
    for(auto i:v[x])
    {
        if(i!=f[x])
        {
            f[i]=x,dfs1(i);
            siz[x]+=siz[i];
            if(siz[son[x]]<siz[i]) son[x]=i;
        }
    }
}
void dfs2(int x,int tp)
{
    top[x]=tp,id[x]=++cnt,rk[cnt]=x;
    if(son[x]) dfs2(son[x],tp);
    for(auto i:v[x]) if(i!=f[x]&&i!=son[x]) dfs2(i,i);
}
struct node
{
    int l,r;
    mutable int v;
    node(int L,int R=-1,int V=0){l=L,r=R,v=V;}
    bool operator < (const node &x) const
    {
        return l<x.l;
    }
};
set <node> s;
set <node> :: iterator split(int pos)
{
    auto it=s.lower_bound(node(pos));
    if(it!=s.end()&&it->l==pos) return it;
    it--;
    if(it->r<pos) return s.end();
    int L=it->l,R=it->r,V=it->v;
    s.erase(it);
    s.insert(node(L,pos-1,V));
    return s.insert(node(pos,R,V)).fi;
}
void modify(int l,int r,int k)
{
    auto it2=split(r+1),it1=split(l);
    s.erase(it1,it2);
    s.insert(node(l,r,k));
}
int query(int l,int r)
{
    int ans=0;
    auto it2=split(r+1),it1=split(l);
    for(;it1!=it2;it1++) last==it1->v?0:ans++,last=it1->v;
    return ans;
}
void update(int x,int y,int k)
{
    while(top[x]!=top[y])
    {
        if(dep[top[x]]<dep[top[y]]) swap(x,y);
        modify(id[top[x]],id[x],k);
        x=f[top[x]];
    }
    if(dep[x]>dep[y]) swap(x,y);
    modify(id[x],id[y],k);
}
int query1(int x,int y)
{
    int ans=0;
    while(top[x]!=top[y])
    {
        if(dep[top[x]]<dep[top[y]]) swap(x,y);
        ans+=query(id[top[x]],id[x]);
        auto it2=split(id[top[x]]),it1=split(id[f[top[x]]]);
        if(it1->v==it2->v) ans--;
        last=-1;
        x=f[top[x]];
    }
    if(dep[x]>dep[y]) swap(x,y);
    ans+=query(id[x],id[y]);
    return ans;
}
int main()
{
    int n,m;
    read(n),read(m);
    for(re int i=1;i<=n;i++) read(val[i]);
    for(re int i=1,x,y;i<n;i++) read(x),read(y),v[x].pb(y),v[y].pb(x);
    dfs1(rt),dfs2(rt,rt);
    for(re int i=1;i<=n;i++) s.insert(node(id[i],id[i],val[i]));
    s.insert(node(n+1,n+1,0)); char op;
    for(re int i=1,l,r,k;i<=m;i++)
    {
        scanf(" %c %d %d",&op,&l,&r);
        if(op==‘Q‘) printf("%d\n",query1(l,r));
        else if(op==‘C‘) read(k),update(l,r,k);
        last=-1;
    }
    return 0;
}