HYSBZ 2243 染色 树链剖分 点上剖分

#include <bits/stdc++.h>
#define sf scanf
#define pf printf

using namespace std;
const int maxn = 100000 + 5;
int n,m;
//边集
struct Edge{
int v,pre,c;
}Es[maxn * 2];
int head[maxn],TOT_EDGE;
void INIT_EDGE(){memset(head,-1,sizeof head);TOT_EDGE = 0;}
void ADD_EDGE(int u,int v,int c){
Es[TOT_EDGE].v = v;
Es[TOT_EDGE].c = c;
Es[TOT_EDGE].pre = head[u];
head[u] = TOT_EDGE++;
}

//树链剖分
int son[maxn],fa[maxn],dep[maxn],size[maxn];
int top[maxn],tid[maxn],rank[maxn],SegSize;
int DFS(int rt){
dep[rt] = dep[fa[rt]] + 1;
son[rt] = 0;
size[rt] = 1;
for(int i = head[rt];~i;i = Es[i].pre){
int v = Es[i].v;
if(v != fa[rt]){
fa[v] = rt;
size[rt] += DFS(v);
if(size[son[rt]] < size[v]) son[rt] = v;
}
}
return size[rt];
}
void Split(int rt,int tp){
top[rt] = tp;
tid[rt] = ++SegSize;
rank[tid[rt]] = rt;
if(son[rt]){
Split(son[rt],tp);
for(int i = head[rt];~i;i = Es[i].pre){
int v = Es[i].v;
if(v != fa[rt] && v != son[rt]) Split(v,v);
}
}
}
void TreeLineSplit(){
dep[0] = 0;
size[0] = 0;
fa[1] = 0;
DFS(1);
SegSize = 0;
Split(1,1);
}

//线段树
#define lson rt << 1 , l , mid
#define rson rt << 1 | 1,mid + 1,r
int LEFT[maxn << 2],RIGHT[maxn << 2],CNT[maxn << 2],num[maxn];
int cover[maxn << 2];
void PushUp(int rt){
LEFT[rt] = LEFT[rt << 1];
RIGHT[rt] = RIGHT[rt << 1 | 1];
CNT[rt] = CNT[rt << 1] + CNT[rt << 1 | 1] - (RIGHT[rt << 1] == LEFT[rt << 1 | 1]);
}
void PushDown(int rt){
if(cover[rt] != -1){
LEFT[rt << 1] = RIGHT[rt << 1] = cover[rt];
LEFT[rt << 1 | 1] = RIGHT[rt << 1 | 1] = cover[rt];
CNT[rt << 1] = CNT[rt << 1 | 1] = 1;
cover[rt << 1] = cover[rt << 1 | 1] = cover[rt];
cover[rt] = -1;
}
}
void SegTreeBuild(int rt,int l,int r){
cover[rt] = -1;
if(l == r){
LEFT[rt] = RIGHT[rt] = num[l];
CNT[rt] = 1;
return;
}
int mid = l + r >> 1;
SegTreeBuild(lson);SegTreeBuild(rson);
PushUp(rt);
}
int LEFT_COLOR,RIGHT_COLOR;
int Query(int rt,int l,int r,int L,int R){
if(L <= l && R >= r){
if(LEFT_COLOR == -1) LEFT_COLOR = LEFT[rt];
RIGHT_COLOR = RIGHT[rt];
return CNT[rt];
}
PushDown(rt);
int mid = l + r >> 1,ret = 0,f = 0;
if(L <= mid) ret += Query(lson,L,R),f++;
if(R > mid) ret += Query(rson,L,R),f++;
if(f == 2 && RIGHT[rt << 1] == LEFT[rt << 1 | 1]) ret--;
PushUp(rt);
return ret;
}
void Update(int rt,int l,int r,int L,int R,int c){
if(L <= l && R >= r){
cover[rt] = c;
LEFT[rt] = RIGHT[rt] = c;
CNT[rt] = 1;
return;
}
PushDown(rt);
int mid = l + r >> 1;
if(L <= mid) Update(lson,L,R,c);
if(R > mid) Update(rson,L,R,c);
PushUp(rt);
}
int INIT_COLOR[maxn];
void GetNumAr(){
for(int i = 1;i <= n;++i){
num[tid[i]] = INIT_COLOR[i];
}
}

//题目操作
int Ans_Query(int u,int v){
int ret = 0;
int preU = -1,preV = -1;
while(top[u] != top[v]){
LEFT_COLOR = -1;
if(dep[top[u]] > dep[top[v]]){
ret += Query(1,1,n,tid[top[u]],tid[u]);
if(RIGHT_COLOR == preU){
ret--;
}
preU = LEFT_COLOR;
u = fa[top[u]];
}
else{
ret += Query(1,1,n,tid[top[v]],tid[v]);
if(RIGHT_COLOR == preV){
ret--;
}
preV = LEFT_COLOR;
v = fa[top[v]];
}
}
LEFT_COLOR = -1;
if(dep[u] < dep[v]){
ret += Query(1,1,n,tid[u],tid[v]);
if(RIGHT_COLOR == preV) ret--;
if(LEFT_COLOR == preU) ret--;
}
else{
ret += Query(1,1,n,tid[v],tid[u]);
if(RIGHT_COLOR == preU) ret--;
if(LEFT_COLOR == preV) ret--;

}
return ret;
}

void CHANGE(int u,int v,int c){
while(top[u] != top[v]){
if(dep[top[v]] < dep[top[u]]) swap(u,v);
Update(1,1,n,tid[top[v]],tid[v],c);
v = fa[top[v]];
}
if(dep[v] < dep[u]) swap(u,v);
Update(1,1,n,tid[u],tid[v],c);
}
char op[10];
int main(){
//    freopen("rand.txt","r",stdin);
while( ~sf("%d%d",&n,&m) ){
INIT_EDGE();
for(int i = 1;i <= n;++i) sf("%d",&INIT_COLOR[i]);
for(int i = 1;i < n;++i){
int u,v;sf("%d%d",&u,&v);
ADD_EDGE(u,v,0);
ADD_EDGE(v,u,0);
}
TreeLineSplit();
GetNumAr();
SegTreeBuild(1,1,n);
LEFT_COLOR = -1;
while(m--){
int x,y,z;
sf("%s",op);
if(op[0] == 'Q'){
sf("%d%d",&x,&y);
pf("%d\n",Ans_Query(x,y));
}
else{
sf("%d%d%d",&x,&y,&z);
CHANGE(x,y,z);
}
}
}
}