2243: [SDOI2011]染色(树链剖分+线段树)

2243: [SDOI2011]染色

Time Limit: 20 Sec  Memory Limit: 512 MB
Submit: 8400  Solved: 3150
[Submit][Status][Discuss]

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<iostream>
#include<cstdio>
#include<cstring>

#define N 1000007

using namespace std;
int head[N],siz[N],fa[N],deep[N];
int pos[N],top[N],col[N],sum[N];
int n,m,ans,cnt,num,tot;
struct edge{int u,v,net;}e[N<<];
struct tree{int l,r,sum,lc,rc,flag;}tr[N<<];

inline int read()
{
    int x=,f=;char c=getchar();
    while(c>''||c<''){if(c=='-')f=-;c=getchar();}
    while(c>=''&&c<=''){x=x*+c-'';c=getchar();}
    return x*f;
}

inline void add(int u,int v)
{
    e[++cnt].v=v;e[cnt].net=head[u];head[u]=cnt;
}

inline void pushup(int k)
{
    if(tr[k<<|].lc==tr[k<<].rc)
    tr[k].sum=tr[k<<].sum+tr[k<<|].sum-;
    else tr[k].sum=tr[k<<].sum+tr[k<<|].sum;
    tr[k].lc=tr[k<<].lc;tr[k].rc=tr[k<<|].rc;
}

inline void pushdown(int k)
{
    if(tr[k].l==tr[k].r) return;
    tr[k<<].sum=tr[k<<|].sum=;
    tr[k<<].lc=tr[k<<].rc=tr[k].flag;
    tr[k<<|].lc=tr[k<<|].rc=tr[k].flag;
    tr[k<<].flag=tr[k<<|].flag=tr[k].flag;
    tr[k].flag=;
}

void build(int k,int l,int r)
{
    tr[k].l=l;tr[k].r=r;
    if(l==r)
    {
        tr[k].sum=;
        tr[k].lc=tr[k].rc=sum[++num];
        return;
    }
    int mid=(l+r)>>;
    build(k<<,l,mid);build(k<<|,mid+,r);
    pushup(k);
}

void change(int k,int l,int r,int c)
{
    if(tr[k].l==l && tr[k].r==r)
    {
        tr[k].sum=;
        tr[k].flag=tr[k].lc=tr[k].rc=c;
        return;
    }
    if(tr[k].flag) pushdown(k);
    int mid=tr[k].l+tr[k].r>>;
    if(r<=mid) change(k<<,l,r,c);
    else if(l>mid) change(k<<|,l,r,c);
    else change(k<<,l,mid,c),change(k<<|,mid+,r,c);
    pushup(k);
}

int query(int k,int l,int r)
{
    if(tr[k].l==l && tr[k].r==r) return tr[k].sum;
    if(tr[k].flag) pushdown(k);
    pushup(k);
    int mid=tr[k].l+tr[k].r>>;
    if(r<=mid) return query(k<<,l,r);
    else if(l>mid) return query(k<<|,l,r);
    else
    {
        int ld=query(k<<,l,mid);
        int rd=query(k<<|,mid+,r);
        if(tr[k<<].rc==tr[k<<|].lc) return ld+rd-;
        else return ld+rd;
    }
}

void dfs1(int u)
{
    siz[u]=;
    for(int i=head[u];i;i=e[i].net)
    {
        int v=e[i].v;
        if(v==fa[u])continue;
        fa[v]=u;deep[v]=deep[u]+;
        dfs1(v);siz[u]+=siz[v];
    }
}

void dfs2(int u,int Top)
{
    top[u]=Top;int k=;
    pos[u]=++tot;sum[pos[u]]=col[u];
    for(int i=head[u];i;i=e[i].net)
    {
        int v=e[i].v;
        if(v==fa[u]) continue;
        if(siz[v]>siz[k]) k=v;
    }if(k!=) dfs2(k,Top);
    else return;
    for(int i=head[u];i;i=e[i].net)
    {
        int v=e[i].v;
        if(v==fa[u] || v==k)continue;
        dfs2(v,v);
    }return;
}

int querycolor(int k,int pos)
{
    if(tr[k].l==tr[k].r && tr[k].l==pos) return tr[k].lc;
    if(tr[k].flag) pushdown(k);
    pushup(k);
    int mid=tr[k].l+tr[k].r>>;
    if(pos<=mid) return querycolor(k<<,pos);
    else return querycolor(k<<|,pos);
}

int solvequery(int x,int y)
{
    ans=;
    while(top[x]!=top[y])
    {
        if(deep[top[x]]<deep[top[y]])swap(x,y);
        ans+=query(,pos[top[x]],pos[x]);
        if(querycolor(,pos[top[x]])==querycolor(,pos[fa[top[x]]])) ans--;
        x=fa[top[x]];
    }
    ans+=query(,min(pos[x],pos[y]),max(pos[x],pos[y]));
    return ans;
}

void solvechange(int x,int y,int c)
{
    while(top[x]!=top[y])
    {
        if(deep[top[x]]<deep[top[y]])swap(x,y);
        change(,pos[top[x]],pos[x],c);
        x=fa[top[x]];
    }
    change(,min(pos[x],pos[y]),max(pos[x],pos[y]),c);
}

int main()
{
    int x,y,z;char ch[];
    n=read();m=read();
    for(int i=;i<=n;i++) col[i]=read();
    for(int i=;i<n;i++)
    {
        x=read();y=read();
        add(x,y);add(y,x);
    }
    dfs1();dfs2(,);
    build(,,n);
    for(int i=;i<=m;i++)
    {
        scanf("%s",ch);
        if(ch[]=='Q')
        {
            x=read();y=read();
            printf("%d\n",solvequery(x,y));
        }
        else
        {
            x=read();y=read();z=read();
            solvechange(x,y,z);
        }
    }
    return ;
}