题解:
树剖,线段树维护区间颜色段数
记录两端点的颜色,做到O(1)合并
问题:
非递归建树实现
#include<iostream> #include<cstdio> #include<cstring> using namespace std; const int maxn=100009; int n,m; int incolor[maxn]; int cntedge; int head[maxn]; int to[maxn<<1],nex[maxn<<1]; int addedge(int x,int y){ nex[++cntedge]=head[x]; to[cntedge]=y; head[x]=cntedge; } int father[maxn],depth[maxn],siz[maxn]; int heavyson[maxn],top[maxn],idx[maxn],ref[maxn]; int dfs1(int now,int fa){ father[now]=fa; depth[now]=depth[fa]+1; siz[now]=1; for(int i=head[now];i;i=nex[i]){ if(to[i]==fa)continue; dfs1(to[i],now); siz[now]+=siz[to[i]]; if(siz[to[i]]>siz[heavyson[now]]){ heavyson[now]=to[i]; } } } int temp; int dfs2(int now,int toppoint){ top[now]=toppoint; idx[now]=(++temp); ref[temp]=now; if(!heavyson[now])return 0; dfs2(heavyson[now],toppoint); for(int i=head[now];i;i=nex[i]){ if(to[i]==father[now])continue; if(to[i]==heavyson[now])continue; dfs2(to[i],to[i]); } } struct SegmentTree{ int l,r; int lcolor,rcolor; int setcolor,sum; }tree[maxn<<2]; int pushup(int now){ tree[now].lcolor=tree[now<<1].lcolor; tree[now].rcolor=tree[now<<1|1].rcolor; if(tree[now<<1].rcolor==tree[now<<1|1].lcolor){ tree[now].sum=tree[now<<1].sum+tree[now<<1|1].sum-1; }else{ tree[now].sum=tree[now<<1].sum+tree[now<<1|1].sum; } } int pushdown(int now){ if(tree[now].setcolor!=-1){ tree[now<<1].setcolor=tree[now<<1].lcolor=tree[now<<1].rcolor=tree[now].setcolor; tree[now<<1|1].setcolor=tree[now<<1|1].lcolor=tree[now<<1|1].rcolor=tree[now].setcolor; tree[now<<1].sum=tree[now<<1|1].sum=1; tree[now].setcolor=-1; } } int Build(int now,int l,int r){ tree[now].l=l;tree[now].r=r; tree[now].setcolor=-1; if(l==r){ tree[now].lcolor=tree[now].rcolor=incolor[ref[l]]; tree[now].sum=1; return 0; } int mid=(l+r)>>1; Build(now<<1,l,mid); Build(now<<1|1,mid+1,r); pushup(now); } int Updata(int now,int ll,int rr,int x){ if(tree[now].l>=ll&&tree[now].r<=rr){ tree[now].setcolor=tree[now].lcolor=tree[now].rcolor=x; tree[now].sum=1; return 0; } int mid=(tree[now].l+tree[now].r)>>1; pushdown(now); if(ll<=mid)Updata(now<<1,ll,rr,x); if(rr>mid)Updata(now<<1|1,ll,rr,x); pushup(now); } int Querysum(int now,int ll,int rr){ if(tree[now].l>=ll&&tree[now].r<=rr){ return tree[now].sum; } int mid=(tree[now].l+tree[now].r)>>1; pushdown(now); if(rr<=mid)return Querysum(now<<1,ll,rr); else if(ll>mid)return Querysum(now<<1|1,ll,rr); else if(tree[now<<1].rcolor==tree[now<<1|1].lcolor)return Querysum(now<<1,ll,rr)+Querysum(now<<1|1,ll,rr)-1; else return Querysum(now<<1,ll,rr)+Querysum(now<<1|1,ll,rr); } int Querycolor(int now,int p){ if(tree[now].l==tree[now].r){ return tree[now].lcolor; } int mid=(tree[now].l+tree[now].r)>>1; pushdown(now); if(p<=mid)return Querycolor(now<<1,p); else return Querycolor(now<<1|1,p); } int Change(int u,int v,int x){ int tu=top[u]; int tv=top[v]; while(tu!=tv){ if(depth[tu]<depth[tv]){ swap(tu,tv);swap(u,v); } Updata(1,idx[tu],idx[u],x); u=father[tu];tu=top[u]; } if(depth[u]>depth[v])swap(u,v); Updata(1,idx[u],idx[v],x); } int Getans(int u,int v){ int ret=0; int tu=top[u]; int tv=top[v]; while(tu!=tv){ if(depth[tu]<depth[tv]){ swap(tu,tv);swap(u,v); } ret=ret+Querysum(1,idx[tu],idx[u]); u=father[tu]; if(Querycolor(1,idx[u])==Querycolor(1,idx[tu]))--ret; tu=top[u]; } if(depth[u]>depth[v])swap(u,v); ret=ret+Querysum(1,idx[u],idx[v]); return ret; } int minit(){ temp=cntedge=0; memset(heavyson,0,sizeof(heavyson)); memset(head,0,sizeof(head)); } int main(){ scanf("%d%d",&n,&m); minit(); for(int i=1;i<=n;++i)scanf("%d",&incolor[i]); for(int i=1;i<=n-1;++i){ int x,y; scanf("%d%d",&x,&y); addedge(x,y);addedge(y,x); } dfs1(1,0); dfs2(1,1); Build(1,1,n); char opty[4]; while(m--){ int x,y,z; scanf("%s",opty); if(opty[0]=='C'){ scanf("%d%d%d",&x,&y,&z); Change(x,y,z); }else{ scanf("%d%d",&x,&y); printf("%d\n",Getans(x,y)); } } return 0; }