bzoj2243 [SDOI2011]染色

#include<cstdio>
const int M = 1001010;
const int N = 300010;
int dp,p
,pre
,tt
,size
,f
,gf
,go
,deep
,color
;
int id
,n,m,i,a,b,c,tot,e
;
int l[M],r[M],lc[M],rc[M],s[M],v[M];
void link(int x,int y)
{
dp++;pre[dp]=p[x];p[x]=dp;tt[dp]=y;
}
struct O{
int ans,lc,rc;
void clear()
{
ans=0;lc=0;rc=0;
}
void set(int a,int b,int c)
{
ans=a;lc=b;rc=c;
}
};
void dfs(int x,int fa)
{
int i=p[x];
size[x]=1;
deep[x]=deep[fa]+1;
while (i)
{
if (tt[i]!=fa)
{
f[tt[i]]=x;
dfs(tt[i],x);
if (size[tt[i]]>size[go[x]]) go[x]=tt[i];
size[x]+=size[tt[i]];
}
i=pre[i];
}
}
void gao(int x,int Fa,int fa)
{
int i=p[x];
gf[x]=Fa;id[x]=++tot;e[tot]=color[x];
if (go[x]) gao(go[x],Fa,x);
while (i)
{
if ((tt[i]!=fa)&&(tt[i]!=go[x]))
gao(tt[i],tt[i],x);
i=pre[i];
}
}
void updata(int x)
{
lc[x]=lc[2*x];
rc[x]=rc[2*x+1];
s[x]=s[2*x]+s[2*x+1];
if (rc[2*x]==lc[2*x+1]) s[x]--;
}
void build(int x,int a,int b)
{
int m;
l[x]=a;r[x]=b;v[x]=-1;
if (b-a>1)
{
m=(a+b)>>1;
build(2*x,a,m);
build(2*x+1,m,b);
updata(x);
}
else
{
lc[x]=e[b];
rc[x]=e[b];
s[x]=1;
}
}
void clean(int x)
{
if (v[x]>=0)
{
s[x]=1;
lc[x]=rc[x]=v[x];
v[2*x]=v[x];
v[2*x+1]=v[x];
v[x]=-1;
}
}
void change(int x,int a,int b,int c)
{
int m;
clean(x);
if ((a<=l[x])&&(r[x]<=b))
{
v[x]=c;
return;
}
m=(l[x]+r[x])>>1;
if (a<m) change(2*x,a,b,c);
if (m<b) change(2*x+1,a,b,c);
clean(2*x);clean(2*x+1);
updata(x);
}
O query(int x,int a,int b)
{
int m;
O tmp;tmp.clear();
clean(x);
if ((a<=l[x])&&(r[x]<=b))
{
tmp.set(s[x],lc[x],rc[x]);
return tmp;
}
m=(l[x]+r[x])>>1;
int L=-1,R=-1,A=-1;
if (a<m)
{
tmp=query(2*x,a,b);
L=tmp.lc;
R=tmp.rc;
A=tmp.ans;
}
if (m<b)
{
tmp=query(2*x+1,a,b);
if (L==-1)
{
L=tmp.lc;
R=tmp.rc;
A=tmp.ans;
}
else
{
A+=tmp.ans;
if (R==tmp.lc) A--;
R=tmp.rc;
}
}
tmp.set(A,L,R);return tmp;
}
void gai(int a,int b,int c)
{
while (1)
{
if (gf[a]==gf[b])
{
if (deep[a]<deep[b]) a^=b^=a^=b;
change(1,id[b]-1,id[a],c);
return;
}
else
{
if (deep[gf[a]]<deep[gf[b]]) a^=b^=a^=b;
change(1,id[gf[a]]-1,id[a],c);
a=f[gf[a]];
}
}
}
int wen(int a,int b)
{
int ac=-1,bc=-1,ans=0;
while (1)
{
if (gf[a]==gf[b])
{
if (deep[a]<deep[b]) a^=b^=a^=b,ac^=bc^=ac^=bc;
O tmp=query(1,id[b]-1,id[a]);
ans+=tmp.ans;
if (tmp.rc==ac) ans--;
if (tmp.lc==bc) ans--;
return ans;
}
else
{
if (deep[gf[a]]<deep[gf[b]]) a^=b^=a^=b,ac^=bc^=ac^=bc;
O tmp=query(1,id[gf[a]]-1,id[a]);
ans+=tmp.ans;
if (tmp.rc==ac) ans--;ac=tmp.lc;
a=f[gf[a]];
}
}
}
int main()
{
scanf("%d%d",&n,&m);
for (i=1;i<=n;i++)
scanf("%d",&color[i]);
for (i=1;i<n;i++)
{
scanf("%d%d",&a,&b);
link(a,b);
link(b,a);
}
dfs(1,0);
gao(1,1,0);
build(1,0,n);
for (i=1;i<=m;i++)
{
char typ;
scanf(" %c",&typ);
if (typ=='Q')
{
scanf("%d%d",&a,&b);
printf("%d\n",wen(a,b));
}
else
{
scanf("%d%d%d",&a,&b,&c);
gai(a,b,c);
}
}
}