P2486 [SDOI2011]染色

題目描述

輸入輸出格式

對於每個詢問操作，輸出一行答案。

輸入輸出樣例

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

3
1
2

說明

思路：裸的樹鏈剖分。

優化一下time

#include<cstdio>
#include
 
<cstring>
#include<iostream>
#include<algorithm>
#define N 100001
using namespace std;
int n,m,sz,tot;
int to[N*2],net[N*2],col[N],head[N];
int deep[N],size[N],dad[N],top[N],id[N];
struct nond{
    int l,r,flag;
    int l_col,r_col,sum;
}tree[N*4];
void add(int u,int v){
    to[
 
++tot]=v;net[tot]=head[u];head[u]=tot;
    to[++tot]=u;net[tot]=head[v];head[v]=tot;
}
void up(int now){ 
    tree[now].l_col=tree[now*2].l_col;
    tree[now].r_col=tree[now*2+1].r_col;
    tree[now].sum=tree[now*2].sum+tree[now*2+1].sum-(tree[now*2].r_col==tree[now*2+1].l_col);
} 
void build(int
 
 now,int l,int r){
    tree[now].l=l;tree[now].r=r;
    if(tree[now].l==tree[now].r)    return ;
    int mid=(tree[now].l+tree[now].r)/2;
    build(now*2,l,mid);
    build(now*2+1,mid+1,r);
    up(now);
}
void change(int now,int t,int x){
    if(tree[now].l==tree[now].r){
        tree[now].sum=1;
        tree[now].l_col=tree[now].r_col=tree[now].flag=x;
        return;
    }
    int mid=(tree[now].l+tree[now].r)/2;
    if(t<=mid)    change(now*2,t,x);
    else if(t>mid)    change(now*2+1,t,x);
    up(now);
}
void down(int now){
    tree[now*2].sum=tree[now*2+1].sum=1;
    tree[now*2].flag=tree[now*2+1].flag=tree[now].flag;
    tree[now*2].l_col=tree[now*2+1].l_col=tree[now].flag;
    tree[now*2].r_col=tree[now*2+1].r_col=tree[now].flag;
    tree[now].flag=0;
}
void change_many(int now,int pol,int por,int a){
    if(tree[now].l==pol&&tree[now].r==por){
        tree[now].flag=tree[now].l_col=tree[now].r_col=a;
        tree[now].sum=1;
        return;
    }
    if(tree[now].flag)    down(now);
    int mid=(tree[now].l+tree[now].r)/2;
    if(por<=mid)    change_many(now*2,pol,por,a);
    else if(pol>mid)    change_many(now*2+1,pol,por,a);
    else{ change_many(now*2,pol,mid,a); change_many(now*2+1,mid+1,por,a); }
    up(now);
}
int query(int now,int t){
    if(tree[now].l==tree[now].r)
        return tree[now].flag;
    if(tree[now].flag)    down(now);
    int mid=(tree[now].l+tree[now].r)/2;
    if(t<=mid)    query(now*2,t);
    else query(now*2+1,t);
}
int query_many(int now,int pol,int por){
    if(tree[now].l==pol&&tree[now].r==por)
        return tree[now].sum;
    if(tree[now].flag)    down(now);
    int mid=(tree[now].l+tree[now].r)/2;
    if(por<=mid)    return query_many(now*2,pol,por);
    else if(pol>mid)    return query_many(now*2+1,pol,por);
    else{
        int t=query_many(now*2,pol,mid)+query_many(now*2+1,mid+1,por);
        if(tree[now*2].r_col==tree[now*2+1].l_col)    t--;
        return t;
    }
} 
int dfs(int now){
    size[now]=1;
    deep[now]=deep[dad[now]]+1;
    for(int i=head[now];i;i=net[i])
        if(dad[now]!=to[i]){
            dad[to[i]]=now;
            dfs(to[i]);
            size[now]+=size[to[i]];
        }
}
int dfs1(int x){
    int t=0;id[x]=++sz;
    if(!top[x])    top[x]=x;
    change(1,id[x],col[x]);
    for(int i=head[x];i;i=net[i])
        if(to[i]!=dad[x]&&size[to[i]]>size[t])
            t=to[i];
    if(t){
        top[t]=top[x];
        dfs1(t);
    }
    for(int i=head[x];i;i=net[i])
        if(to[i]!=dad[x]&&to[i]!=t)
            dfs1(to[i]);
}
int schange(int x,int y,int a){
    for(;top[x]!=top[y];){
        if(deep[top[x]]<deep[top[y]])    swap(x,y);
        change_many(1,id[top[x]],id[x],a);
        x=dad[top[x]];
    }
    if(id[x]>id[y])    swap(x,y);
    change_many(1,id[x],id[y],a);
}
void squery(int x,int y){
    int ans=0;
    while(top[x]!=top[y]){
        if(deep[top[x]]<deep[top[y]])    swap(x,y);
        ans+=query_many(1,id[top[x]],id[x]);
        if(query(1,id[dad[top[x]]])==query(1,id[top[x]]))    ans--;
        x=dad[top[x]];
    }
    if(id[x]>id[y])    swap(x,y);
    ans+=query_many(1,id[x],id[y]);
    cout<<ans<<endl;
}
int main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)    scanf("%d",&col[i]);
    for(int i=1;i<n;i++){
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
    }
    build(1,1,n);dfs(1);dfs1(1);
    for(int i=1;i<=m;i++){
        char ch;int a,b,c;
        cin>>ch;scanf("%d%d",&a,&b);
        if(ch==‘C‘){ scanf("%d",&c);schange(a,b,c); }
        else    squery(a,b);
    }
}

洛谷 P2486 [SDOI2011]染色