bzoj2243 [SDOI2011]染色 树链剖分

Description

---

给定一棵有n个节点的无根树和m个操作，操作有2类：
1、将节点a到节点b路径上所有点都染成颜色c；
2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），
如“112221”由3段组成：“11”、“222”和“1”。
请你写一个程序依次完成这m个操作。

“C a b c”表示这是一个染色操作，把节点a到节点b路径上所有点（包括a和b）都染成颜色c；
“Q a b”表示这是一个询问操作，询问节点a到节点b（包括a和b）路径上的颜色段数量。

HINT

---

数N<=10^5，操作数M<=10^5，所有的颜色C为整数且在[0, 10^9]之间。

Solution

---

很自然地就想到了树链剖分，线段树维护左边第一个颜色，右边第一个颜色，以及这一段颜色的数量

考虑到在查询重边的时候会有相邻颜色相同的一段，那么我们每做完一条重边就判断它的父亲是否和重边中较浅的点颜色相同

Code

---

#include <stdio.h>
#include <string.h>
#define rep(i,st,ed) for (int i=st;i<=ed;++i)
#define fill(x,t) memset(x,t,sizeof(x))
const int N=200005;
const int E=500005;
struct edge{int x,y,next;}e[E];
int lazy[N<<2],tot[N<<2],lc[N<<2],rc[N<<2];
int size[N],pos[N],dep[N],bl[N],fa[N],cnt=0;
int ls[N],edCnt=0;
int c[N];
int n,m;
int read() {
    int x=0,v=1; char ch=getchar();
    for (;ch<'0'||ch>'9';v=(ch=='-')?(-1):(v),ch=getchar());
    for (;ch<='9'&&ch>='0';x=x*10+ch-'0',ch=getchar());
    return x*v;
}
void addEdge(int x,int y) {
    e[++edCnt]=(edge){x,y,ls[x]}; ls[x]=edCnt;
    e[++edCnt]=(edge){y,x,ls[y]}; ls[y]=edCnt;
}
void swap(int &x,int &y) {x^=y; y^=x; x^=y;}
void buildTree(int now,int tl,int tr) {
    if (tl==tr) {return ;}
    int mid=(tl+tr)>>1;
    buildTree(now<<1,tl,mid);
    buildTree(now<<1|1,mid+1,tr);
}
void dfs1(int now) {
    size[now]=1;
    for (int i=ls[now];i;i=e[i].next) {
        if (e[i].y==fa[now]) {continue;}
        dep[e[i].y]=dep[now]+1;
        fa[e[i].y]=now;
        dfs1(e[i].y);
        size[now]+=size[e[i].y];
    }
}
void dfs2(int now,int up) {
    pos[now]=++cnt;
    bl[now]=up;
    int mx=0;
    for (int i=ls[now];i;i=e[i].next) {
        if (dep[e[i].y]>dep[now]&&size[e[i].y]>size[mx]) {mx=e[i].y;}
    }
    if (!mx) {return ;}
    dfs2(mx,up);
    for (int i=ls[now];i;i=e[i].next) {
        if (dep[e[i].y]<dep[now]||e[i].y==mx) {continue;}
        dfs2(e[i].y,e[i].y);
    }
}
void push_down(int now) {
    if (!lazy[now]) {return ;}
    lazy[now<<1]=lazy[now<<1|1]=lazy[now];
    lc[now<<1]=rc[now<<1]=lc[now<<1|1]=rc[now<<1|1]=lazy[now];
    tot[now<<1]=tot[now<<1|1]=1;
    lazy[now]=0;
}
void push_up(int now) {
    tot[now]=tot[now<<1]+tot[now<<1|1];
    if (rc[now<<1]==lc[now<<1|1]) tot[now]-=1;
    lc[now]=lc[now<<1];
    rc[now]=rc[now<<1|1];
}
void modify(int now,int tl,int tr,int l,int r,int v) {
    push_down(now);
    if (tl==l&&tr==r) {
        lazy[now]=v;
        tot[now]=1;
        lc[now]=rc[now]=v;
        return ;
    }
    int mid=(tl+tr)>>1;
    if (r<=mid) {
        modify(now<<1,tl,mid,l,r,v);
    } else if (l>mid) {
        modify(now<<1|1,mid+1,tr,l,r,v);
    } else {
        modify(now<<1,tl,mid,l,mid,v);
        modify(now<<1|1,mid+1,tr,mid+1,r,v);
    }
    push_up(now);
}
int query(int now,int tl,int tr,int l,int r) {
    push_down(now);
    if (tl==l&&tr==r) {return tot[now];}
    int mid=(tl+tr)>>1;
    if (r<=mid) {return query(now<<1,tl,mid,l,r);}
    else if (l>mid) {return query(now<<1|1,mid+1,tr,l,r);}
    else {
        int qx=query(now<<1,tl,mid,l,mid);
        int qy=query(now<<1|1,mid+1,tr,mid+1,r);
        int ret=qx+qy;
        if (rc[now<<1]==lc[now<<1|1]) ret-=1;
        return ret;
    }
}
int query_color(int now,int tl,int tr,int pos) {
    push_down(now);
    if (tl==tr) {return lc[now];}
    int mid=(tl+tr)>>1;
    if (pos<=mid) {return query_color(now<<1,tl,mid,pos);}
    else {return query_color(now<<1|1,mid+1,tr,pos);}
}
int get_ans(int x,int y) {
    int ret=0;
    while (bl[x]!=bl[y]) {
        if (dep[bl[x]]<dep[bl[y]]) {swap(x,y);}
        ret+=query(1,1,n,pos[bl[x]],pos[x]);
        int color1=query_color(1,1,n,pos[bl[x]]);
        int color2=query_color(1,1,n,pos[fa[bl[x]]]);
        if (color1==color2) {ret-=1;}
        x=fa[bl[x]];
    }
    if (dep[x]>dep[y]) {swap(x,y);}
    ret+=query(1,1,n,pos[x],pos[y]);
    return ret;
}
void solve(int x,int y,int z) {
    while (bl[x]!=bl[y]) {
        if (dep[bl[x]]<dep[bl[y]]) {swap(x,y);}
        modify(1,1,n,pos[bl[x]],pos[x],z);
        x=fa[bl[x]];
    }
    if (dep[x]>dep[y]) {swap(x,y);}
    modify(1,1,n,pos[x],pos[y],z);
}
int main(void) {
    n=read();
    m=read();
    rep(i,1,n) {c[i]=read();}
    rep(i,2,n) {
        int x=read();
        int y=read();
        addEdge(x,y);
    }
    dfs1(1);
    dfs2(1,1);
    rep(i,1,n) {modify(1,1,n,pos[i],pos[i],c[i]);}
    while (m--) {
        char opt[5]; scanf("%s",opt);
        if (opt[0]=='Q') {
            int x=read();
            int y=read();
            int prt=get_ans(x,y);
            printf("%d\n",prt);
        } else if (opt[0]=='C') {
            int x=read();
            int y=read();
            int z=read();
            solve(x,y,z);
        }
    }
    return 0;
}