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
对于每个询问操作,输出一行答案。
树剖之后对线段树的每个节点记录颜色数、左端点颜色和右端点颜色,方便合并。
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
int fa[100010],son[100010],dep[100010],top[100010],pos[100010],size[100010],
num[1000010],tag[1000010],lc[1000010],rc[1000010],
fir[100010],ne[200010],to[200010],clr[100010],
m,n,tot;
void add(int num,int u,int v)
{ne[num]=fir[u];fir[u]=num;to[num]=v;
}
void init()
{int i,x,y;scanf("%d%d",&n,&m);for (i=1;i<=n;i++)scanf("%d",&clr[i]);for (i=1;i<n;i++){scanf("%d%d",&x,&y);add(i*2,x,y);add(i*2+1,y,x);}
}
void down(int p)
{if (tag[p]>=0){tag[p*2]=tag[p*2+1]=lc[p]=rc[p]=tag[p];num[p]=1;tag[p]=-1;}
}
void up(int p)
{down(p);down(p*2);down(p*2+1);lc[p]=lc[p*2];rc[p]=rc[p*2+1];num[p]=num[p*2]+num[p*2+1]-(bool)(rc[p*2]==lc[p*2+1]);
}
void modi(int p,int L,int R,int l,int r,int x)
{down(p);if (L==l&&R==r){tag[p]=x;return;}int mid=(L+R)/2;if (r<=mid) modi(p*2,L,mid,l,r,x);else{if (l>=mid+1) modi(p*2+1,mid+1,R,l,r,x);else{modi(p*2,L,mid,l,mid,x);modi(p*2+1,mid+1,R,mid+1,r,x);}}up(p);
}
int qry(int p,int L,int R,int l,int r,int &ll,int &rr)
{down(p);if (L==l&&R==r){ll=lc[p];rr=rc[p];return num[p];}int mid=(L+R)/2;if (r<=mid) return qry(p*2,L,mid,l,r,ll,rr);if (l>=mid+1) return qry(p*2+1,mid+1,R,l,r,ll,rr);int a1,l1,r1,a2,l2,r2;a1=qry(p*2,L,mid,l,mid,l1,r1);a2=qry(p*2+1,mid+1,R,mid+1,r,l2,r2);ll=l1;rr=r2;return a1+a2-(bool)(r1==l2);
}
void dfs1(int u,int f)
{int i,v;size[u]=1;for (i=fir[u];i;i=ne[i])if ((v=to[i])!=f){dep[v]=dep[u]+1;fa[v]=u;dfs1(v,u);size[u]+=size[v];if (son[u]==0||size[v]>size[son[u]])son[u]=v;}
}
void dfs2(int u)
{int i,v;pos[u]=++tot;if (!son[u]) return;top[son[u]]=top[u];dfs2(son[u]);for (i=fir[u];i;i=ne[i])if ((v=to[i])!=son[u]&&v!=fa[u]){top[v]=v;dfs2(v);}
}
void pre()
{int i;dep[1]=1;dfs1(1,-1);top[1]=1;dfs2(1);memset(tag,-1,sizeof(tag));for (i=1;i<=n;i++)modi(1,1,tot,pos[i],pos[i],clr[i]);
}
void modify(int u,int v,int x)
{int f1,f2;while ((f1=top[u])!=(f2=top[v])){if (dep[f1]<dep[f2]){swap(f1,f2);swap(u,v);}modi(1,1,tot,pos[f1],pos[u],x);u=fa[f1];}if (dep[u]<dep[v])swap(u,v);modi(1,1,tot,pos[v],pos[u],x);
}
int query(int u,int v)
{int l1=-1,l2=-1,ret=0,f1,f2,tl,tr,temp;while ((f1=top[u])!=(f2=top[v]))if (dep[f1]<dep[f2]){temp=qry(1,1,tot,pos[f2],pos[v],tl,tr);ret+=temp-(bool)(l2==tr);l2=tl;v=fa[f2];}else{temp=qry(1,1,tot,pos[f1],pos[u],tl,tr);ret+=temp-(bool)(l1==tr);l1=tl;u=fa[f1];}if (dep[u]<dep[v]){temp=qry(1,1,tot,pos[u],pos[v],tl,tr);ret+=temp-(bool)(l1==tl)-(bool)(l2==tr);}else{temp=qry(1,1,tot,pos[v],pos[u],tl,tr);ret+=temp-(bool)(l2==tl)-(bool)(l1==tr);}return ret;
}
int main()
{char s[5];int x,y,z;init();pre();while (m--){scanf("%s",s);if (s[0]=='C'){scanf("%d%d%d",&x,&y,&z);modify(x,y,z);}else{scanf("%d%d",&x,&y);printf("%d\n",query(x,y));}}
}