【BZOJ2243】[SDOI2011]染色 树链剖分+线段树

【BZOJ2243】[SDOI2011]染色

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]之间。

题解：树链剖分+线段树的模板题，注意在区间的合并时要判断相邻的端点颜色是否相同。

#include <stdio.h>
#include <string.h>
#include <iostream>
#define lson x<<1
#define rson x<<1|1
using namespace std;
const int maxn=100010;
int n,m,cnt,tot;
int fa[maxn],deep[maxn],size[maxn],son[maxn],top[maxn];
int to[maxn<<1],next[maxn<<1],head[maxn];
int p[maxn],s[maxn<<2],sa[maxn<<2],sb[maxn<<2],u[maxn],v[maxn],tag[maxn<<2];
char str[10];
int readin()
{
    int ret=0,sig=1;    char gc;
    while(gc<'0'||gc>'9')    sig=(gc=='-')?-1:1,gc=getchar();
    while(gc>='0'&&gc<='9')    ret=ret*10+gc-'0',gc=getchar();
    return ret*sig;
}
void add(int a,int b)
{
    to[cnt]=b;
    next[cnt]=head[a];
    head[a]=cnt++;
}
void dfs1(int x)
{
    size[x]=1;
    for(int i=head[x];i!=-1;i=next[i])
    {
        if(to[i]!=fa[x])
        {
            deep[to[i]]=deep[x]+1;
            fa[to[i]]=x;
            dfs1(to[i]);
            size[x]+=size[to[i]];
            if(size[to[i]]>size[son[x]])    son[x]=to[i];
        }
    }
}
void dfs2(int x,int tp)
{
    top[x]=tp;
    p[x]=++tot;
    u[p[x]]=v[x];
    if(son[x])    dfs2(son[x],tp);
    for(int i=head[x];i!=-1;i=next[i])
        if(to[i]!=fa[x]&&to[i]!=son[x])
            dfs2(to[i],to[i]);
}
void pushup(int x)
{
    s[x]=s[lson]+s[rson];
    sa[x]=sa[lson],sb[x]=sb[rson];
    if(sb[lson]==sa[rson])    s[x]--;
}
void pushdown(int x)
{
    if(tag[x])
    {
        sa[lson]=sa[rson]=sb[lson]=sb[rson]=tag[lson]=tag[rson]=tag[x];
        s[lson]=s[rson]=1,tag[x]=0;
    }
}
void build(int l,int r,int x)
{
    if(l==r)
    {
        s[x]=1,sa[x]=sb[x]=u[l];
        return ;
    }
    int mid=l+r>>1;
    build(l,mid,lson),build(mid+1,r,rson);
    pushup(x);
}
void updata(int l,int r,int x,int a,int b,int c)
{
    if(a<=l&&r<=b)
    {
        s[x]=1,sa[x]=sb[x]=tag[x]=c;
        return ;
    }
    pushdown(x);
    int mid=l+r>>1;
    if(b<=mid)    updata(l,mid,lson,a,b,c);
    else if(a>mid)    updata(mid+1,r,rson,a,b,c);
    else    updata(l,mid,lson,a,b,c),updata(mid+1,r,rson,a,b,c);
    pushup(x);
}
void change(int x,int y)
{
    int c=readin();
    while(top[x]!=top[y])
    {
        if(deep[top[x]]<deep[top[y]])    swap(x,y);
        updata(1,n,1,p[top[x]],p[x],c);
        x=fa[top[x]];
    }
    if(deep[x]>deep[y])    swap(x,y);
    updata(1,n,1,p[x],p[y],c);
}
int query(int l,int r,int x,int a,int b)
{
    if(a<=l&&r<=b)    return s[x];
    pushdown(x);
    int mid=l+r>>1;
    if(b<=mid)    return query(l,mid,lson,a,b);
    if(a>mid)    return query(mid+1,r,rson,a,b);
    return query(l,mid,lson,a,b)+query(mid+1,r,rson,a,b)-(sb[lson]==sa[rson]);
}
int getcolor(int l,int r,int x,int y)
{
    if(l==r)    return sa[x];
    pushdown(x);
    int mid=l+r>>1;
    if(y<=mid)    return getcolor(l,mid,lson,y);
    if(y>mid)    return getcolor(mid+1,r,rson,y);
}
int getans(int x,int y)
{
    int ans=0;
    while(top[x]!=top[y])
    {
        if(deep[top[x]]<deep[top[y]])    swap(x,y);
        ans+=query(1,n,1,p[top[x]],p[x]);
        if(getcolor(1,n,1,p[top[x]])==getcolor(1,n,1,p[fa[top[x]]]))    ans--;
            //判断端点的颜色是否相同
        x=fa[top[x]];
    }
    if(deep[x]>deep[y])    swap(x,y);
    ans+=query(1,n,1,p[x],p[y]);
    printf("%d\n",ans);
}
int main()
{
    n=readin(),m=readin();
    int i,a,b;
    memset(head,-1,sizeof(head));
    for(i=1;i<=n;i++)    v[i]=readin();
    for(i=1;i<n;i++)
    {
        a=readin(),b=readin();
        add(a,b),add(b,a);
    }
    deep[1]=1;
    dfs1(1);
    dfs2(1,1);
    build(1,n,1);
    for(i=1;i<=m;i++)
    {
        scanf("%s",str);
        a=readin(),b=readin();
        switch(str[0])
        {
            case 'C':change(a,b);    break;
            case 'Q':getans(a,b);    break;
        }
    }
    return 0;
}