Description
给定一棵有$n$个节点的无根树和$m$个操作,操作有$2$类:
1.将节点$a$到节点$b$路径上所有点都染成颜色$c$;
2.询问节点$a$到节点$b$路径上的颜色段数量(连续相同颜色被认为是同一段),如$"112221"$由$3$段组成:$"11","222"和"1"$.
请你写一个程序依次完成这$m$个操作.
Input
第一行包含$2$个整数$n$和$m$,分别表示节点数和操作数.
第二行包含$n$个正整数表示$n$个节点的初始颜色.
下面$n-1$行每行包含两个整数$x$和$y$,表示$x$和$y$之间有一条无向边.
下面$m$行每行描述一个操作:
$”C\;a\;b\;c”$表示这是一个染色操作,把节点$a$到节点$b$路径上所有点(包括$a$和$b$)都染成颜色$c$;
$”Q\;a\;b”$表示这是一个询问操作,询问节点$a$到节点$b$(包括$a$和$b$)路径上的颜色段数量.
Output
对于每个询问操作,输出一行答案.
Sample Input
6 5
2 2 1 2 1 1
1 2
1 3
2 4
2 5
2 6
Q 3 5
C 2 1 1
Q 3 5
C 5 1 2
Q 3 5
Sample Output
3
1
2
HINT
$N,M\;\leq\;10^5,c\;\in\;z,c\;\in\;[0, 10^9]$.
Solution
树链剖分+线段树.
每次上移的时候判断相邻两段端点是否同色即可.
#include<cmath>
#include<ctime>
#include<queue>
#include<stack>
#include<cstdio>
#include<vector>
#include<cstring>
#include<cstdlib>
#include<iostream>
#include<algorithm>
#define N 100005
#define M 300005
using namespace std;
struct linetree{
int l,r,lc,rc,cnt;
}lt[M];
struct graph{
int nxt,to;
}e[M];
char c[2];
int g[N],a[N],w[N],n,m,u,v,k,l,cnt;
int f[N],p[N],dep[N],siz[N],son[N],top[N];
inline void addedge(int x,int y){
e[++cnt].nxt=g[x];g[x]=cnt;e[cnt].to=y;
}
inline void dfs1(int u){
int m=0;siz[u]=1;
for(int i=g[u];i;i=e[i].nxt)
if(!dep[e[i].to]){
f[e[i].to]=u;
dep[e[i].to]=dep[u]+1;
dfs1(e[i].to);
siz[u]+=siz[e[i].to];
if(siz[e[i].to]>m){
son[u]=e[i].to;
m=siz[e[i].to];
}
}
}
inline void dfs2(int u,int tp){
top[u]=tp;p[u]=++cnt;w[cnt]=a[u];
if(son[u]) dfs2(son[u],tp);
for(int i=g[u];i;i=e[i].nxt)
if(f[u]!=e[i].to&&e[i].to!=son[u])
dfs2(e[i].to,e[i].to);
}
inline void build(int u,int l,int r){
lt[u].l=l;lt[u].r=r;
if(lt[u].l<lt[u].r){
int lef=u<<1,rig=u<<1|1;
int mid=(lt[u].l+lt[u].r)>>1;
build(lef,l,mid);build(rig,mid+1,r);
lt[u].cnt=lt[lef].cnt+lt[rig].cnt;
if(lt[lef].rc==lt[rig].lc) --lt[u].cnt;
lt[u].lc=lt[lef].lc;lt[u].rc=lt[rig].rc;
}
else{
lt[u].cnt=1;lt[u].lc=lt[u].rc=w[lt[u].l];
}
}
inline void cover(int u,int l,int r,int k){
if(lt[u].l>=l&<[u].r<=r){
lt[u].lc=lt[u].rc=k;lt[u].cnt=1;
}
else if(lt[u].l<lt[u].r){
int lef=u<<1,rig=u<<1|1;
int mid=(lt[u].l+lt[u].r)>>1;
if(lt[u].cnt==1){
lt[lef].cnt=lt[rig].cnt=1;
lt[lef].lc=lt[lef].rc=lt[rig].lc=lt[rig].rc=lt[u].lc;
}
if(l<=mid) cover(lef,l,r,k);
if(r>mid) cover(rig,l,r,k);
lt[u].cnt=lt[lef].cnt+lt[rig].cnt;
if(lt[lef].rc==lt[rig].lc) --lt[u].cnt;
lt[u].lc=lt[lef].lc;lt[u].rc=lt[rig].rc;
}
}
inline int ask(int u,int l,int r){
if(lt[u].l>=l&<[u].r<=r)
return lt[u].cnt;
if(lt[u].l<lt[u].r){
int lef=u<<1,rig=u<<1|1,ret=0;
int mid=(lt[u].l+lt[u].r)>>1;
if(lt[u].cnt==1){
lt[lef].cnt=lt[rig].cnt=1;
lt[lef].lc=lt[rig].lc=lt[lef].rc=lt[rig].rc=lt[u].lc;
}
if(l<=mid) ret+=ask(lef,l,r);
if(r>mid) ret+=ask(rig,l,r);
if(l<=mid&&r>mid&<[lef].rc==lt[rig].lc) --ret;
return ret;
}
}
inline int col(int u,int x){
if(lt[u].l==x) return lt[u].lc;
if(lt[u].r==x) return lt[u].rc;
if(lt[u].l<lt[u].r){
int lef=u<<1,rig=u<<1|1;
int mid=(lt[u].l+lt[u].r)>>1;
if(lt[u].cnt==1){
lt[lef].cnt=lt[rig].cnt=1;
lt[lef].lc=lt[rig].lc=lt[lef].rc=lt[rig].rc=lt[u].lc;
}
if(x<=mid) return col(lef,x);
return col(rig,x);
}
}
inline void cov(int x,int y,int k){
int t;
while(top[x]!=top[y]){
if(dep[top[x]]<dep[top[y]]){
t=x;x=y;y=t;
}
cover(1,p[top[x]],p[x],k);
x=f[top[x]];
}
if(p[x]>p[y]){
t=x;x=y;y=t;
}
cover(1,p[x],p[y],k);
}
inline int q(int x,int y){
int ret=0,t;
while(top[x]!=top[y]){
if(dep[top[x]]<dep[top[y]]){
t=x;x=y;y=t;
}
ret+=ask(1,p[top[x]],p[x]);
if(col(1,p[top[x]])==col(1,p[f[top[x]]])) --ret;
x=f[top[x]];
}
if(p[x]>p[y]){
t=x;x=y;y=t;
}
ret+=ask(1,p[x],p[y]);
return ret;
}
inline int lca(int x,int y){
int t;
while(top[x]!=top[y]){
if(dep[top[x]]<dep[top[y]]){
t=x;x=y;y=t;
}
x=f[top[x]];
}
if(p[x]>p[y]){
t=x;x=y;y=t;
}
return x;
}
inline void Aireen(){
scanf("%d%d",&n,&m);
for(int i=1;i<=n;++i)
scanf("%d",&a[i]);
for(int i=1,x,y;i<n;++i){
scanf("%d%d",&x,&y);
addedge(x,y);addedge(y,x);
}
dep[1]=1;dfs1(1);
cnt=0;dfs2(1,1);
build(1,1,n);
while(m--){
scanf("%s%d%d",c,&u,&v);
if(c[0]=='C'){
scanf("%d",&k);cov(u,v,k);
}
else{
l=lca(u,v);
printf("%d\n",q(u,v));
}
}
}
int main(){
freopen("color.in","r",stdin);
freopen("color.out","w",stdout);
Aireen();
fclose(stdin);
fclose(stdout);
return 0;
}