[树链剖分] [bzoj2243] [SDOI2011]染色

Description

给定一棵有n个节点的无根树和m个操作，操作有2类：

1、将节点a到节点b路径上所有点都染成颜色c；

2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），

如“112221”由3段组成：“11”、“222”和“1”。

请你写一个程序依次完成这m个操作。

Input

第一行包含2个整数n和m，分别表示节点数和操作数；

第二行包含n个正整数表示n个节点的初始颜色

下面 行每行包含两个整数x和y，表示x和y之间有一条无向边。

下面 行每行描述一个操作：

“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<=10^5，操作数M<=10^5，所有的颜色C为整数且在[0, 10^9]之间。

题目一看就知道是树剖，就是写起来有点烦

注意线段树询问的时候和树剖的时候都要判断边界颜色是否一样

#include<cstdio>
#include<iostream>
using namespace std;

const int MaxN=100005;

int N,M;
int Tot,Head[MaxN],Next[MaxN<<1],To[MaxN<<1];
int Fa[MaxN],Dep[MaxN],Size[MaxN],Son[MaxN];
int Index,Top[MaxN],ID[MaxN],Point[MaxN];
int Col[MaxN];

struct Node{
int l,r;
int lcol,rcol,cnt;
int tag;
};

struct Ans{
int cnt,lcol,rcol;
};

#define L (Seg[cur].l)
#define R (Seg[cur].r)
#define Mid (Seg[cur].l+Seg[cur].r>>1)
#define LCol (Seg[cur].lcol)
#define RCol (Seg[cur].rcol)
#define Cnt (Seg[cur].cnt)
#define Tag (Seg[cur].tag)
#define LCH (cur<<1)
#define RCH (cur<<1|1)

struct Segment_Tree{
Node Seg[MaxN<<2];

void PushUp(int cur){
LCol=Seg[LCH].lcol,RCol=Seg[RCH].rcol;
Cnt=Seg[LCH].cnt+Seg[RCH].cnt-(Seg[LCH].rcol==Seg[RCH].lcol);
}

void Build(int cur,int l,int r){
L=l,R=r;
if(L==R){
LCol=RCol=Col[Point[l]];
Cnt=1;
return;
}
Build(LCH,L,Mid),Build(RCH,Mid+1,R);
PushUp(cur);
}

void PushDown(int cur){
if(Tag){
Seg[LCH].lcol=Seg[LCH].rcol=Seg[RCH].lcol=Seg[RCH].rcol=Tag;
Seg[LCH].cnt=Seg[RCH].cnt=1;
Seg[LCH].tag=Seg[RCH].tag=Tag;
Tag=0;
}
}

void Update(int cur,int l,int r,int v){
if(L>=l&&R<=r){
LCol=RCol=v,Cnt=1,Tag=v;
return;
}
PushDown(cur);
if(Mid>=l)
Update(LCH,l,r,v);
if(Mid<r)
Update(RCH,l,r,v);
PushUp(cur);
}

int Get_Col(int cur,int pos){
if(L==R)
return LCol;
PushDown(cur);
if(Mid>=pos)
return Get_Col(LCH,pos);
return Get_Col(RCH,pos);
}

Ans Get_Cnt(int cur,int l,int r){
if(L>=l&&R<=r)
return (Ans){Cnt,LCol,RCol};
PushDown(cur);
Ans a=(Ans){0,0,0},b=(Ans){0,0,0};
if(Mid>=l)
a=Get_Cnt(LCH,l,r);
if(Mid<r)
b=Get_Cnt(RCH,l,r);
return (Ans){a.cnt+b.cnt-(Mid>=l&&Mid<r)*(a.rcol==b.lcol),a.lcol,b.rcol};
}
}Seg_Tree;

int Input(){
int s=0;
char ch;
while(ch=getchar(),ch<'0'||ch>'9');
while(ch>='0'&&ch<='9'){
s=s*10+ch-'0';
ch=getchar();
}
return s;
}

void Add_Edge(int u,int v){
Next[++Tot]=Head[u];
To[Tot]=v;
Head[u]=Tot;
}

void DFS1(int u,int fa){
int i,v;
Size[u]=1;
for(i=Head[u];i;i=Next[i])
if((v=To[i])!=fa){
Fa[v]=u,Dep[v]=Dep[u]+1;
DFS1(v,u);
Size[u]+=Size[v];
Son[u]=(Size[Son[u]]>Size[v]?Son[u]:v);
}
}

void DFS2(int u){
int i,v;
Point[ID[u]=++Index]=u;
if(v=Son[u])
Top[v]=Top[u],DFS2(v);
for(i=Head[u];i;i=Next[i])
if((v=To[i])!=Fa[u]&&v!=Son[u])
Top[v]=v,DFS2(v);
}

void Update(int u,int v,int c){
int r1=Top[u],r2=Top[v];
while(r1!=r2){
if(Dep[r1]<Dep[r2])
swap(u,v),swap(r1,r2);
Seg_Tree.Update(1,ID[r1],ID[u],c);
u=Fa[r1],r1=Top[u];
}
if(Dep[u]>Dep[v])
swap(u,v);
Seg_Tree.Update(1,ID[u],ID[v],c);
}

void Query(int u,int v){
int r1=Top[u],r2=Top[v],cnt=0;
while(r1!=r2){
if(Dep[r1]<Dep[r2])
swap(u,v),swap(r1,r2);
cnt+=Seg_Tree.Get_Cnt(1,ID[r1],ID[u]).cnt;
cnt-=(Fa[r1]&&Seg_Tree.Get_Col(1,ID[r1])==Seg_Tree.Get_Col(1,ID[Fa[r1]]));
u=Fa[r1],r1=Top[u];
}
if(Dep[u]>Dep[v])
swap(u,v);
cnt+=Seg_Tree.Get_Cnt(1,ID[u],ID[v]).cnt;
printf("%d\n",cnt);
}

int main(){
int i,u,v,c;
char opt[5];
N=Input(),M=Input();
for(i=1;i<=N;i++)
Col[i]=Input();
for(i=1;i<N;i++){
u=Input(),v=Input();
Add_Edge(u,v);
Add_Edge(v,u);
}
DFS1(1,0);
Top[1]=1,DFS2(1);
Seg_Tree.Build(1,1,N);
while(M--){
scanf("%s",opt),u=Input(),v=Input();
if(*opt=='C'){
c=Input();
Update(u,v,c);
}
else Query(u,v);
}
return 0;
}