给定一颗有n个节点的无根树和m个操作,操作有2类:
1.将节点a到节点b路径上的所有点都染成颜色a;
2.询问节点a到节点b路径上的颜色段数量(连续相同颜色被认为是同一段),如“112221”由3段组成:“11”、“222”、“1”。
思路:冷静分析.jpg。等我什么时候重新理解了这份几年前写的代码再好好详画一下思路吧
#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
using namespace std;
const int MAXN=100005;
struct Edge{int to,next;}edg[MAXN*3];
struct Tree{int l,r,v,l_,r_,laz;}tre[MAXN*4];
int fa[MAXN],son[MAXN],num[MAXN],head[MAXN],e[MAXN][2],top[MAXN];
int dep[MAXN],p[MAXN],tot,m,n,pos,fp[MAXN],color[MAXN];
void adde(int u,int v){
edg[++tot].to=v;edg[tot].next=head[u];head[u]=tot;
}
void dfs1(int u,int pre,int d){
dep[u]=d;
fa[u]=pre;
num[u]=1;
for(int i=head[u];i!=-1;i=edg[i].next){
int v=edg[i].to;
if(v!=fa[u]){
dfs1(v,u,d+1);
num[u]+=num[v];
if(son[u]==-1||num[son[u]]<num[v])son[u]=v;
}
}
}
void dfs2(int u,int sp){
top[u]=sp;
p[u]=++pos;
fp[pos]=u;
if(son[u]!=-1)dfs2(son[u],sp);else return;
for(int i=head[u];i!=-1;i=edg[i].next){
int v=edg[i].to;
if(v!=son[u]&&v!=fa[u])dfs2(v,v);
}
}
void build(int i,int l,int r){
tre[i].l=l;
tre[i].r=r;
if(l==r)return;
int mid=(l+r)/2;
build(i<<1,l,mid);
build((i<<1)|1,mid+1,r);
}
void push_up(int i){
tre[i].v=tre[i<<1].v+tre[(i<<1)|1].v;
if(tre[i<<1].r_==tre[(i<<1)|1].l_&&tre[i<<1].r_)tre[i].v--;
tre[i].l_=tre[i<<1].l_;
tre[i].r_=tre[(i<<1)|1].r_;
}
void init_tre(int i,int k,int val){
if(tre[i].l==k&&tre[i].r==k){
tre[i].v=1;
tre[i].l_=val;
tre[i].r_=val;
return;
}
int mid=(tre[i].l+tre[i].r)/2;
if(k<=mid)init_tre(i<<1,k,val);
else init_tre((i<<1)|1,k,val);
push_up(i);
}
void paint(int i,int x){
tre[i].laz=x;
tre[i].l_=x;
tre[i].r_=x;
tre[i].v=1;
}
void push_down(int i){
paint(i<<1,tre[i].laz);
paint((i<<1)|1,tre[i].laz);
tre[i].laz=0;
}
void update(int i,int l,int r,int val){
if(tre[i].l>=l&&tre[i].r<=r){
tre[i].laz=tre[i].l_=tre[i].r_=val;
tre[i].v=1;
return;
}
if(tre[i].laz)push_down(i);
int mid=(tre[i].l+tre[i].r)/2;
if(r<=mid)update(i<<1,l,r,val);else
if(l>mid)update((i<<1)|1,l,r,val);
else{
update(i<<1,l,mid,val);
update((i<<1)|1,mid+1,r,val);
}
push_up(i);
}
int find_color(int i,int k){
if(tre[i].l==k&&tre[i].r==k)return tre[i].l_;
if(tre[i].laz)push_down(i);
int mid=(tre[i].l+tre[i].r)/2;
if(k<=mid)return find_color(i<<1,k);else
return find_color((i<<1)|1,k);
}
void modify(int u,int v,int val){
int f1=top[u],f2=top[v];
while (f1!=f2){
if(dep[f1]<dep[f2]){
swap(f1,f2);
swap(u,v);
}
update(1,p[f1],p[u],val);
u=fa[f1];f1=top[u];
}
if(dep[u]>dep[v])swap(u,v);
update(1,p[u],p[v],val);
}
int query(int i,int l,int r){
if(tre[i].l==l&&tre[i].r==r)return tre[i].v;
if(tre[i].laz)push_down(i);
int mid=(tre[i].l+tre[i].r)/2;
if(r<=mid)return query(i<<1,l,r);else
if(l>mid)return query((i<<1)|1,l,r);else{
int t=query(i<<1,l,mid)+query((i<<1)|1,mid+1,r);
if(tre[i<<1].r_==tre[(i<<1)|1].l_&&tre[i<<1].r_)t--;
return t;
}
}
int find(int u,int v){
int f1=top[u],f2=top[v],tmp=0;
while(f1!=f2){
if(dep[f1]<dep[f2]){
swap(f1,f2);
swap(u,v);
}
tmp+=query(1,p[f1],p[u]);
if(find_color(1,p[f1])==find_color(1,p[fa[f1]])&&find_color(1,p[f1]))tmp--;
u=fa[f1];f1=top[u];
}
//if(u==v)return tmp==0?1:tmp;
if(dep[u]>dep[v])swap(u,v);
tmp+=query(1,p[u],p[v]);
return tmp==0?1:tmp;
}
int main(){
memset(head,-1,sizeof(head));
memset(son,-1,sizeof(son));
memset(tre,0,sizeof(tre));
cin>>n>>m;
for(int i=1;i<=n;i++)scanf("%d",&color[i]);
for(int i=1;i<=n-1;i++){
scanf("%d%d",&e[i][0],&e[i][1]);
adde(e[i][0],e[i][1]);
adde(e[i][1],e[i][0]);
}
dfs1(1,0,1);
dfs2(1,1);
build(1,1,pos);
for(int i=1;i<=n;i++)init_tre(1,p[i],color[i]);
char op[10];
int u,v,x;
while(scanf("%s",op)==1){
if(op[0]=='Q'){
scanf("%d%d",&u,&v);
printf("%d\n",find(u,v));
}else{
scanf("%d%d%d",&u,&v,&x);
modify(u,v,x);
}
}
return 0;
}