[SDOI2011] 染色
题目描述
给定一棵 n 个节点的无根树，共有 m 个操作，操作分为两种：
- 将节点 a 到节点 b 的路径上的所有点（包括 a 和 b ）都染成颜色 c 。
- 询问节点 a 到节点 b 的路径上的颜色段数量。
颜色段的定义是极长的连续相同颜色被认为是一段。例如 112221 由三段组成：11、222、1。
输入格式
输入的第一行是用空格隔开的两个整数，分别代表树的节点个数 n 和操作个数 $m$。
第二行有 n 个用空格隔开的整数，第 i 个整数 w_i 代表结点 i 的初始颜色。
第 3 到第 (n + 1) 行，每行两个用空格隔开的整数 u, v ，代表树上存在一条连结节点 u 和节点 v 的边。
第 (n + 2) 到第 (n + m + 1) 行，每行描述一个操作，其格式为：
每行首先有一个字符 op ，代表本次操作的类型。
- 若 op 为
C，则代表本次操作是一次染色操作，在一个空格后有三个用空格隔开的整数 a, b, c ，代表将 a 到 b 的路径上所有点都染成颜色 c 。 - 若 op 为
Q，则代表本次操作是一次查询操作，在一个空格后有两个用空格隔开的整数 a, b ，表示查询 a 到 b 路径上的颜色段数量。
输出格式
对于每次查询操作，输出一行一个整数代表答案。
样例 #1
样例输入 #1
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
样例输出 #1
3
1
2
提示
数据规模与约定
对于 100\% 的数据， 1 \leq n, m \leq 10^5 ， 1 \leq w_i, c \leq 10^9 ， 1 \leq a, b, u, v \leq n ， op 一定为 C 或 Q，保证给出的图是一棵树。
除原数据外，还存在一组不计分的 hack 数据。
代码
调了好久，下午继续调，应该是get_c有问题，样例是能过得
#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=114514;
int ls(int p){return (p<<1);}
int rs(int p){return (p<<1)+1;}
int m,n,a[N],w[N];
int cnt,head[N],dfn;
struct edge{
int next,to;
}e[N*2];
void add_edge(int u,int v){
e[++cnt].to=v;
e[cnt].next=head[u];
head[u]=cnt;
}
struct tree{
int size,top,son,dep,fa,id;
tree(){
memset(this,0,sizeof(tree));
}
}p[N];
struct node{
int lc,rc,sum,lazy;
node(){
memset(this,0,sizeof(node));
}
}sgm[N*4];
void dfs1(int u,int fa){
p[u].fa=fa;
p[u].dep=p[fa].dep+1;
p[u].size=1;
for(int i=head[u];i;i=e[i].next){
int v=e[i].to;
if(v==fa) continue;
dfs1(v,u);
p[u].size+=p[v].size;
if(p[u].son==0||p[v].size>p[p[u].son].size){
p[u].son=v;
}
}
}
void dfs2(int u,int t){
p[u].id=++dfn;
p[u].top=t;
if(!p[u].son) return;
dfs2(p[u].son,t);
for(int i=head[u];i;i=e[i].next){
int v=e[i].to;
if(v==p[u].fa||v==p[u].son) continue;
dfs2(v,v);
}
}
void pushup(int p,int pl,int pr){
int mid=(pl+pr)>>1;
sgm[p].sum=sgm[ls(p)].sum+sgm[rs(p)].sum;
if(sgm[ls(p)].rc==sgm[rs(p)].lc) {sgm[p].sum--;}
sgm[p].lc=sgm[ls(p)].lc;
sgm[p].rc=sgm[rs(p)].rc;
}
void assign(int p,int pl,int pr,int k){//推平
sgm[p].lazy=k;
sgm[p].sum=1;
sgm[p].lc=sgm[p].rc=k;
}
void pushdown(int p,int pl,int pr){
if(sgm[p].lazy){
int mid=(pl+pr)>>1;
assign(ls(p),pl,mid,sgm[p].lazy);
assign(rs(p),mid+1,pr,sgm[p].lazy);
sgm[p].lazy=0;
}
}
int get_c(int p,int pl,int pr,int k){
if(pl==k){return sgm[p].lc;}
if(pr==k){return sgm[p].rc;}
pushdown(p,pl,pr);
int mid=(pl+pr)>>1;
if(k<=mid) get_c(ls(p),pl,mid,k);
else get_c(rs(p),mid+1,pr,k);
}
void build(int p,int pl,int pr){
if(pl==pr){
sgm[p].lc=sgm[p].rc=w[pl];
sgm[p].sum=1;
return;
}
int mid=(pl+pr)>>1;
build(ls(p),pl,mid);
build(rs(p),mid+1,pr);
pushup(p,pl,pr);
}
int query(int p,int pl,int pr,int l,int r){
if(l<=pl&&pr<=r){return sgm[p].sum;}
int res=0,mid=(pl+pr)>>1;
pushdown(p,pl,pr);
if(l<=mid) {res+=query(ls(p),pl,mid,l,r);}//cout<<pl<<" "<<mid<<" "<<res<<endl;
if(mid+1<=r) {res+=query(rs(p),mid+1,pr,l,r);}//cout<<(mid+1)<<" "<<pr<<" "<<res<<endl;
if(l<=mid&&mid+1<=r&&sgm[ls(p)].rc==sgm[rs(p)].lc) {res--;}
return res;
}
void update(int p,int pl,int pr,int l,int r,int k){
if(l<=pl&&pr<=r){assign(p,pl,pr,k);return ;}
int mid=(pl+pr)>>1;
pushdown(p,pl,pr);
if(l<=mid) update(ls(p),pl,mid,l,r,k);
if(mid+1<=r) update(rs(p),mid+1,pr,l,r,k);
pushup(p,pl,pr);
}
void update_c(int x,int y,int k){
while(p[x].top!=p[y].top){
if(p[p[x].top].dep<p[p[y].top].dep) swap(x,y);
update(1,1,n,p[p[x].top].id,p[x].id,k);
x=p[p[x].top].fa;
//cout<<x<<endl;
}
if(p[x].dep<p[y].dep) swap(x,y);
//cout<<"update:"<<p[y].id<<" "<<p[x].id<<" "<<k<<endl;
update(1,1,n,p[y].id,p[x].id,k);
}
int query_c(int x,int y){
int res=0;
while(p[x].top!=p[y].top){
if(p[p[x].top].dep<p[p[y].top].dep) swap(x,y);
res+=query(1,1,n,p[p[x].top].id,p[x].id);
if(get_c(1,1,n,p[p[x].top].id)==get_c(1,1,n,p[p[p[x].top].fa].id)) {
res--;
}
x=p[p[x].top].fa;
}
if(p[x].dep<p[y].dep) swap(x,y);
res+=query(1,1,n,p[y].id,p[x].id);
return res;
}
signed main(){
cin>>n>>m;
for(int i=1;i<=n;i++){
cin>>a[i];
}
for(int i=1;i<=n-1;i++){
int x,y;
cin>>x>>y;
add_edge(x,y);
add_edge(y,x);
}
dfs1(1,1);
dfs2(1,1);
for(int i=1;i<=n;i++){
cout<<p[i].id<<" ";
w[p[i].id]=a[i];
}
cout<<endl;
build(1,1,n);
while(m--){
int x,y,z;
char op;
cin>>op;
if(op=='Q'){
cin>>x>>y;
cout<<query_c(x,y)<<endl;
}
else {
cin>>x>>y>>z;
update_c(x,y,z);
/*
for(int i=1;i<=n;i++){
cout<<get_c(1,1,n,i)<<endl;
}
cout<<endl;
*/
}
}
return 0;
}