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
对于每个询问操作,输出一行答案。
Sample Input
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
Sample Output
3
1
2
HINT
数N<=10^5,操作数M<=10^5,所有的颜色C为整数且在[0, 10^9]之间。
Solution
树链剖分即可,颜色段数还是很好用线段树在区间维护的,就是注意一下树链合并的时候是有方向性的,最后合并的时候要把其中一条链反过来(自己脑补一下)。
代码:
#include<cstdio>
#include<algorithm>
using namespace std;template<typename T>inline void read(T &x){T f=1;char ch=getchar();for(;ch<'0'||ch>'9';ch=getchar())if(ch=='-')f=-1;for(x=0;ch>='0'&&ch<='9';ch=getchar())x=x*10+ch-'0';x*=f;
}const int maxn=100010;
struct edge{int to,nxt;
}e[maxn<<1];
struct Data{int cl,cr,num;Data(){cl=cr=num=0;}friend Data operator+(Data a,Data b){if((!a.num)||(!b.num))return (!a.num)?b:a;Data x;x.cl=a.cl;x.cr=b.cr;x.num=a.num+b.num-(a.cr==b.cl);return x;}
};
struct Seg_Tree{#define lc x<<1#define rc x<<1|1int L[maxn<<2],R[maxn<<2],same[maxn<<2];Data T[maxn<<2];void Build(int x,int *a,int l,int r){same[x]=-1;if((L[x]=l)==(R[x]=r)){T[x].cl=T[x].cr=a[l];return T[x].num=1,void();}int mid=(l+r)>>1;Build(lc,a,l,mid);Build(rc,a,mid+1,r);T[x]=T[lc]+T[rc];}void pushsame(int x,int c){T[x].cl=T[x].cr=same[x]=c;T[x].num=1;}void pushdown(int x){if(same[x]==-1)return;pushsame(lc,same[x]);pushsame(rc,same[x]);same[x]=-1;}void Modify(int x,int l,int r,int c){if(L[x]>=l&&R[x]<=r)return pushsame(x,c),void();int mid=(L[x]+R[x])>>1;pushdown(x);if(l<=mid)Modify(lc,l,r,c);if(r>mid)Modify(rc,l,r,c);T[x]=T[lc]+T[rc];}Data Query(int x,int l,int r){if(L[x]>=l&&R[x]<=r)return T[x];int mid=(L[x]+R[x])>>1;pushdown(x);if(l<=mid&&r>mid)return Query(lc,l,r)+Query(rc,l,r);else if(l<=mid)return Query(lc,l,r);else return Query(rc,l,r);}
}tree;
int n,m,num,head[maxn],col[maxn],dep[maxn];
int fa[maxn],son[maxn],top[maxn],size[maxn],rnk[maxn],pos[maxn],tot;void add(int u,int v){e[++num].to=v;e[num].nxt=head[u];head[u]=num;e[++num].to=u;e[num].nxt=head[v];head[v]=num;
}
void Dfs(int x){size[x]=1;dep[x]=dep[fa[x]]+1;for(int i=head[x];i;i=e[i].nxt)if(e[i].to!=fa[x]){fa[e[i].to]=x;Dfs(e[i].to);size[x]+=size[e[i].to];if(size[e[i].to]>size[son[x]])son[x]=e[i].to;}
}
void Dfs(int x,int tp){top[x]=tp;pos[rnk[x]=++tot]=col[x];if(son[x])Dfs(son[x],tp);for(int i=head[x];i;i=e[i].nxt)if(e[i].to!=fa[x]&&e[i].to!=son[x])Dfs(e[i].to,e[i].to);
}
void Modify(int u,int v,int c){int x=top[u],y=top[v];while(x!=y){if(dep[x]>dep[y]){tree.Modify(1,rnk[x],rnk[u],c);x=top[u=fa[x]];}else{tree.Modify(1,rnk[y],rnk[v],c);y=top[v=fa[y]];}}if(dep[u]<dep[v])tree.Modify(1,rnk[u],rnk[v],c);else tree.Modify(1,rnk[v],rnk[u],c);
}
int Query(int u,int v){int x=top[u],y=top[v];Data ur,vr,ans;while(x!=y){if(dep[x]>dep[y]){ur=tree.Query(1,rnk[x],rnk[u])+ur;x=top[u=fa[x]];}else{vr=tree.Query(1,rnk[y],rnk[v])+vr;y=top[v=fa[y]];}}if(dep[u]<dep[v])vr=tree.Query(1,rnk[u],rnk[v])+vr;else ur=tree.Query(1,rnk[v],rnk[u])+ur;swap(ur.cl,ur.cr);ans=ur+vr;return ans.num;
}int main(){read(n);read(m);for(int i=1;i<=n;i++)read(col[i]);for(int i=1,u,v;i<n;i++){read(u);read(v);add(u,v);}Dfs(1);Dfs(1,1);tree.Build(1,pos,1,tot);while(m--){char opt[5];int u,v,c;scanf("%s",opt);read(u);read(v);if(opt[0]=='C')read(c),Modify(u,v,c);else printf("%d\n",Query(u,v));}return 0;
}