题目大意:
给定一棵有n个节点的无根树和m个操作,操作有2类:
1、将节点a到节点b路径上所有点都染成颜色c;
2、询问节点a到节点b路径上的颜色段数量(连续相同颜色被认为是同一段),
如“112221”由3段组成:“11”、“222”和“1”。
请你写一个程序依次完成这m个操作。
分析:
其实并没有脱离最基础的树链剖分模板。
同样是树链剖分套线段树,线段树内储存当前区间的不同染色段的个数,以及左右端点,每次合并答案时,检查右区间的左端点颜色,和左区间的右端点颜色是否相同,相同则ans-1,
另外,在询问路径时,要多储存两个信息,分别表示u点之前的颜色,和v点之前的颜色,以此来判断是否需要修正ans。
另外,线段树实现时需要用到懒标记,那是线段树的知识了,这里不做过多赘述。
代码量有些大,不过这种思路简单实现复杂的题目,一般来说,尽管代码可能很长,但很难写错,可以很容易一发过,写的时候注意心态,不要着急写错了就能很快写完。
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<vector>
#define SF scanf
#define PF printf
#define MAXN 100010
using namespace std;
struct node{
int val,l,r,tagl,tagr;
node() {}
node(int val1,int ll,int rr):val(val1),l(ll),r(rr) {}
}tree[MAXN*4];
int n,m,q,dcnt,col[MAXN];
int valx[MAXN],fa[MAXN],dep[MAXN],sz[MAXN],son[MAXN],top[MAXN],tid[MAXN],rnk[MAXN];
vector<int> a[MAXN];
void dfs1(int x,int fax,int deep){
fa[x]=fax;
dep[x]=deep;
sz[x]=1;
son[x]=-1;
for(int i=0;i<a[x].size();i++){
if(a[x][i]==fax)
continue;
dfs1(a[x][i],x,deep+1);
sz[x]+=sz[a[x][i]];
if(son[x]==-1||sz[son[x]]<sz[a[x][i]])
son[x]=a[x][i];
}
}
void dfs2(int x,int tp){
top[x]=tp;
tid[x]=++dcnt;
rnk[dcnt]=x;
if(son[x]==-1) return ;
dfs2(son[x],tp);
for(int i=0;i<a[x].size();i++)
if(a[x][i]!=son[x]&&a[x][i]!=fa[x])
dfs2(a[x][i],a[x][i]);
}
void pushup(int x,int tag){
tree[x].val=1;
tree[x].l=tag;
tree[x].r=tag;
tree[x].tagl=tag;
tree[x].tagr=tag;
}
void update(int x){
if(tree[x].tagl!=-1)
pushup(x*2,tree[x].tagl);
if(tree[x].tagr!=-1)
pushup(x*2+1,tree[x].tagr);
tree[x].val=tree[x*2].val+tree[x*2+1].val;
if(tree[x*2].r==tree[x*2+1].l)
tree[x].val--;
tree[x].l=tree[x*2].l;
tree[x].r=tree[x*2+1].r;
}
void build(int l,int r,int id){
if(l==r){
tree[id]=node(1,col[rnk[l]],col[rnk[r]]);
tree[id].tagl=-1;
tree[id].tagr=-1;
return ;
}
int mid=(l+r)>>1;
build(l,mid,id*2);
build(mid+1,r,id*2+1);
tree[id].tagl=-1;
tree[id].tagr=-1;
update(id);
}
node Que(int l,int r,int id,int l1,int r1,int tag){
if(tag!=-1){
tree[id].val=1;
tree[id].l=tag;
tree[id].r=tag;
tree[id].tagl=tag;
tree[id].tagr=tag;
}
if(l1<=l&&r<=r1)
return tree[id];
int mid=(l+r)>>1;
if(mid>=r1){
node x=Que(l,mid,id*2,l1,r1,tree[id].tagl);
tree[id].tagl=-1;
update(id);
return x;
}
if(mid<l1){
node x=Que(mid+1,r,id*2+1,l1,r1,tree[id].tagr);
tree[id].tagr=-1;
update(id);
return x;
}
node x=Que(l,mid,id*2,l1,r1,tree[id].tagl);
node y=Que(mid+1,r,id*2+1,l1,r1,tree[id].tagr);
tree[id].tagl=-1;
tree[id].tagr=-1;
int res=x.val+y.val;
if(x.r==y.l)
res--;
update(id);
return node(res,x.l,y.r);
}
void Change(int l,int r,int id,int l1,int r1,int val,int tag){
if(tag!=-1){
tree[id].val=1;
tree[id].l=tag;
tree[id].r=tag;
tree[id].tagl=tag;
tree[id].tagr=tag;
}
if(l1<=l&&r<=r1){
tree[id].val=1;
tree[id].l=val;
tree[id].r=val;
tree[id].tagl=val;
tree[id].tagr=val;
return ;
}
int mid=(l+r)>>1;
if(mid>=r1){
Change(l,mid,id*2,l1,r1,val,tree[id].tagl);
tree[id].tagl=-1;
update(id);
return ;
}
if(mid<l1){
Change(mid+1,r,id*2+1,l1,r1,val,tree[id].tagr);
tree[id].tagr=-1;
update(id);
return ;
}
Change(l,mid,id*2,l1,r1,val,tree[id].tagl);
Change(mid+1,r,id*2+1,l1,r1,val,tree[id].tagr);
tree[id].tagl=-1;
tree[id].tagr=-1;
update(id);
return ;
}
int Query(int x,int y){
int x1=-1,y1=-1,ans=0;
while(top[x]!=top[y]){
if(dep[top[x]]<dep[top[y]]){
swap(x,y);
swap(x1,y1);
}
node n1=Que(1,n,1,tid[top[x]],tid[x],-1);
ans+=n1.val;
if(n1.r==x1)
ans--;
x1=n1.l;
x=fa[top[x]];
}
if(dep[x]>dep[y]){
swap(x,y);
swap(x1,y1);
}
node n1=Que(1,n,1,tid[x],tid[y],-1);
ans+=n1.val;
if(x1==n1.l)
ans--;
if(y1==n1.r)
ans--;
return ans;
}
void Chan(int x,int y,int val){
while(top[x]!=top[y]){
if(dep[top[x]]<dep[top[y]])
swap(x,y);
Change(1,n,1,tid[top[x]],tid[x],val,-1);
x=fa[top[x]];
}
if(dep[x]>dep[y])
swap(x,y);
Change(1,n,1,tid[x],tid[y],val,-1);
}
int u,v,val;
char s[20];
int main(){
SF("%d%d",&n,&m);
for(int i=1;i<=n;i++)
SF("%d",&col[i]);
for(int i=1;i<n;i++){
SF("%d%d",&u,&v);
a[u].push_back(v);
a[v].push_back(u);
}
dfs1(1,0,1);
dfs2(1,1);
build(1,n,1);
for(int i=1;i<=m;i++){
SF("%s",s);
if(s[0]=='Q'){
SF("%d%d",&u,&v);
PF("%d\n",Query(u,v));
}
else{
SF("%d%d%d",&u,&v,&val);
Chan(u,v,val);
}
}
}