BZOJ 2243 [SDOI2011]染色：树剖【维护路径上颜色段】

题目链接：http://www.lydsy.com/JudgeOnline/problem.php?id=2243

题意：

　　给定一棵有n个节点的无根树和m个操作，操作有2类：

　　1、将节点a到节点b路径上所有点都染成颜色c；

　　2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），如"112221"由3段组成："11"、"222"和"1"。

　　请你写一个程序依次完成这m个操作。

题解：

　　树剖。

　　线段树上维护区间颜色段，每个节点：

　　　　dat[]：当前区间的颜色段数量

　　　　st[]：修改标记

　　　　lc[] and rc[]：当前区间的最左端和最右端颜色

　　　　lson[] and rson[]：左右儿子编号

　　线段树部分见NOI 2007 项链工厂。

　　树剖query中，每次a = par[tp[a]]往上跳之前，先判断getcol(dfsx[tp[a]]) == getcol(dfsx[par[tp[a]]])

　　如果相等，上下两端合并，sum--。

AC Code:

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <vector>
#define MAX_N 1000005

using namespace std;

int n,m;
int tot=0;
int cnt=0;
int c[MAX_N];
int dat[MAX_N];
int st[MAX_N];
int lc[MAX_N];
int rc[MAX_N];
int lson[MAX_N];
int rson[MAX_N];
int par[MAX_N];
int dep[MAX_N];
int siz[MAX_N];
int tp[MAX_N];
int son[MAX_N];
int dfsx[MAX_N];
int idx[MAX_N];
vector<int> edge[MAX_N];

void read()
{
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&c[i]);
    }
    int a,b;
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&a,&b);
        edge[a].push_back(b);
        edge[b].push_back(a);
    }
}

void push_down(int now)
{
    if(st[now]!=-1)
    {
        dat[lson[now]]=dat[rson[now]]=1;
        st[lson[now]]=st[rson[now]]=st[now];
        lc[lson[now]]=rc[lson[now]]=lc[rson[now]]=rc[rson[now]]=st[now];
        st[now]=-1;
    }
}

void push_up(int now)
{
    dat[now]=0;
    lc[now]=lc[lson[now]];
    rc[now]=rc[rson[now]];
    if(lson[now]) dat[now]+=dat[lson[now]];
    if(rson[now]) dat[now]+=dat[rson[now]];
    if(lson[now] && rson[now] && rc[lson[now]]==lc[rson[now]]) dat[now]--;
}

int build(int l,int r)
{
    int rt=++tot;
    dat[rt]=1;
    st[rt]=-1;
    lc[rt]=c[idx[l]];
    rc[rt]=c[idx[r]];
    lson[rt]=rson[rt]=0;
    if(l<r)
    {
        int mid=(l+r)>>1;
        lson[rt]=build(l,mid);
        rson[rt]=build(mid+1,r);
        push_up(rt);
    }
    return rt;
}

void update(int a,int b,int k,int l,int r,int x)
{
    if(a<=l && r<=b)
    {
        dat[k]=1;
        st[k]=lc[k]=rc[k]=x;
        return;
    }
    if(r<a || b<l) return;
    push_down(k);
    int mid=(l+r)>>1;
    update(a,b,lson[k],l,mid,x);
    update(a,b,rson[k],mid+1,r,x);
    push_up(k);
}

int query(int a,int b,int k,int l,int r)
{
    if(a<=l && r<=b) return dat[k];
    if(r<a || b<l) return 0;
    push_down(k);
    int mid=(l+r)>>1;
    int ans=0;
    if(a<=mid) ans+=query(a,b,lson[k],l,mid);
    if(b>mid) ans+=query(a,b,rson[k],mid+1,r);
    if(a<=mid && b>mid && rc[lson[k]]==lc[rson[k]]) ans--;
    return ans;
}

int getcol(int p,int k,int l,int r)
{
    if(l==r) return lc[k];
    push_down(k);
    int mid=(l+r)>>1;
    if(p<=mid) return getcol(p,lson[k],l,mid);
    else return getcol(p,rson[k],mid+1,r);
}

void dfs1(int now,int p,int d)
{
    par[now]=p;
    dep[now]=d;
    siz[now]=1;
    for(int i=0;i<edge[now].size();i++)
    {
        int temp=edge[now][i];
        if(temp!=p)
        {
            dfs1(temp,now,d+1);
            siz[now]+=siz[temp];
        }
    }
}

void dfs2(int now,int anc)
{
    tp[now]=anc;
    son[now]=-1;
    cnt++;
    idx[cnt]=now;
    dfsx[now]=cnt;
    for(int i=0;i<edge[now].size();i++)
    {
        int temp=edge[now][i];
        if((son[now]==-1 || siz[temp]>siz[son[now]]) && temp!=par[now])
        {
            son[now]=temp;
        }
    }
    if(son[now]!=-1) dfs2(son[now],anc);
    for(int i=0;i<edge[now].size();i++)
    {
        int temp=edge[now][i];
        if(temp!=par[now] && temp!=son[now]) dfs2(temp,temp);
    }
}

void update_chain(int a,int b,int x)
{
    while(tp[a]!=tp[b])
    {
        if(dep[tp[a]]<dep[tp[b]]) swap(a,b);
        update(dfsx[tp[a]],dfsx[a],1,1,n,x);
        a=par[tp[a]];
    }
    if(dep[a]<dep[b]) swap(a,b);
    update(dfsx[b],dfsx[a],1,1,n,x);
}

int query_chain(int a,int b)
{
    int sum=0;
    while(tp[a]!=tp[b])
    {
        if(dep[tp[a]]<dep[tp[b]]) swap(a,b);
        sum+=query(dfsx[tp[a]],dfsx[a],1,1,n);
        if(getcol(dfsx[tp[a]],1,1,n)==getcol(dfsx[par[tp[a]]],1,1,n)) sum--;
        a=par[tp[a]];
    }
    if(dep[a]<dep[b]) swap(a,b);
    sum+=query(dfsx[b],dfsx[a],1,1,n);
    return sum;
}

void work()
{
    dfs1(1,-1,0);
    dfs2(1,1);
    build(1,n);
    char s[16];
    int a,b,c;
    while(m--)
    {
        scanf("%s",s);
        if(s[0]=='Q')
        {
            scanf("%d%d",&a,&b);
            printf("%d
",query_chain(a,b));
        }
        else
        {
            scanf("%d%d%d",&a,&b,&c);
            update_chain(a,b,c);
        }
    }
}

int main()
{
    read();
    work();
}