树剖裸题,复习一波。。一开始还看错题目了以为要维护一段区间的颜色数量。。还以为要上主席树。。
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<iostream>
#include<queue>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
using namespace std;
const int N=5e5+5;
int n,m,cnt;
int val[N],head[N],next[N],go[N],fa[N],pos[N];
int a[N],tot;
int siz[N],son[N],top[N],dep[N];
struct node
{
int l,r,sum,lc,rc,lazy;
}t[N*10];
inline void add(int x,int y)
{
go[++tot]=y;
next[tot]=head[x];
head[x]=tot;
}
inline void dfs1(int x,int f)
{
siz[x]=1;
for(int i=head[x];i;i=next[i])
{
int v=go[i];
if (v!=f)
{
dep[v]=dep[x]+1;
fa[v]=x;
dfs1(v,x);
siz[x]+=siz[v];
if (!son[x]||siz[son[x]]<siz[v])
son[x]=v;
}
}
}
int tim;
inline void dfs2(int x,int chain)
{
int k=0;
pos[x]=++tim;
top[x]=chain;
for(int i=head[x];i;i=next[i])
{
int v=go[i];
if (v!=fa[x])
{
if (siz[v]>siz[k])
k=v;
}
}
if (!k)return ;
dfs2(k,chain);
for(int i=head[x];i;i=next[i])
{
int v=go[i];
if (v!=fa[x]&&v!=k)
dfs2(v,v);
}
}
inline void build(int x,int l,int r)
{
t[x].l=l,t[x].r=r;
if (l==r)return ;
int mid=(l+r)>>1;
build(x*2,l,mid);
build(x*2+1,mid+1,r);
}
inline void pushlazy(int x)
{
if (t[x].lazy==0||t[x].l==t[x].r)return;
int v=t[x].lazy;
t[x].lazy=0;
t[x<<1].sum=t[x<<1|1].sum=1;
t[x<<1].lc=t[x<<1].rc=t[x<<1|1].lc=t[x<<1|1].rc=t[x<<1].lazy=t[x<<1|1].lazy=v;
}
inline void queryc(int x,int l,int r,int v)
{
if (t[x].l==l&&t[x].r==r)
{
t[x].sum=1;
t[x].lc=t[x].rc=t[x].lazy=v;
return;
}
pushlazy(x);
int mid=(t[x].l+t[x].r)>>1;
if (r<=mid)queryc(x<<1,l,r,v);
else if (l>mid)queryc(x<<1|1,l,r,v);
else
{
queryc(x<<1,l,mid,v);
queryc(x<<1|1,mid+1,r,v);
}
t[x].lc=t[x<<1].lc;
t[x].rc=t[x<<1|1].rc;
t[x].sum=t[x<<1].sum+t[x<<1|1].sum;
if (t[x<<1].rc==t[x<<1|1].lc)t[x].sum--;
}
inline void solvec(int x,int y,int z)
{
while (top[x]!=top[y])
{
if (dep[top[x]]<dep[top[y]])swap(x,y);
queryc(1,pos[top[x]],pos[x],z);
x=fa[top[x]];
}
if (dep[x]>dep[y])swap(x,y);
queryc(1,pos[x],pos[y],z);
}
inline int queryq(int x,int l,int r)
{
pushlazy(x);
if (t[x].l==l&&t[x].r==r)return t[x].sum;
int mid=(t[x].l+t[x].r)>>1;
if (r<=mid)return queryq(x<<1,l,r);
else if (l>mid)return queryq(x<<1|1,l,r);
else
{
int sum=queryq(x<<1,l,mid)+queryq(x<<1|1,mid+1,r);
if (t[x<<1].rc==t[x<<1|1].lc)sum--;
return sum;
}
}
inline int find(int x,int pos)
{
pushlazy(x);
if (t[x].sum==1)return t[x].lc;
int mid=(t[x].l+t[x].r)>>1;
if (pos<=mid)return find(x<<1,pos);
else find(x<<1|1,pos);
}
inline int solveq(int x,int y)
{
int sum=0;
while (top[x]!=top[y])
{
if (dep[top[x]]<dep[top[y]])swap(x,y);
sum+=queryq(1,pos[top[x]],pos[x]);
if (find(1,pos[top[x]])==find(1,pos[fa[top[x]]]))sum--;//注意合并的时候相邻颜色相同sum-1
x=fa[top[x]];
}
if (dep[x]>dep[y])swap(x,y);
sum+=queryq(1,pos[x],pos[y]);
return sum;
}
int main()
{
scanf("%d%d",&n,&m);
fo(i,1,n)scanf("%d",&a[i]);
fo(i,1,n-1)
{
int x,y;
scanf("%d%d",&x,&y);
add(x,y);
add(y,x);
}
dep[1]=1;
dfs1(1,0);
dfs2(1,1);
build(1,1,n);
char ch[2];
fo(i,1,n)
queryc(1,pos[i],pos[i],a[i]);
fo(i,1,m)
{
scanf("%s",ch);
if (ch[0]=='C')
{
int x,y,z;
scanf("%d%d%d",&x,&y,&z);
solvec(x,y,z);
}
else
{
int x,y;
scanf("%d%d",&x,&y);
printf("%d\n",solveq(x,y));
}
}
}