BZOJ2243 [SDOI2011]染色 题解&代码

题意：
给定一棵有n个节点的树和m个操作，操作有：
C a b c 将树上a到b路径上所有点都染成颜色c；
Q a b 询问树上a到b路径上的颜色段数量（连续相同颜色是同一段）
思路：
树上的路径！树链剖分！
可惜智障了…没想到怎么维护颜色段【妈的这么简单的维护当时居然不会
树剖划分一下树，然后线段树维护每一段的最左lc[]最右rc[]和不同颜色色段数量和sum[]，查询的时候关于判断树中被切开的段的左右端是否一样还是需要谨慎，最后我参考了一下黄学长的处理思路来着

最后：颜色可能有0

/**************************************************************
    Problem: 2243
    User: Rainbow6174
    Language: C++
    Result: Accepted
    Time:9056 ms
    Memory:20356 kb
****************************************************************/

#include<iostream>
#include<vector>
#include<cstdio>
#include<cstring>
#define lson (o<<1)
#define rson ((o<<1)|1)
using namespace std;
const int maxn = 100005;
int n,m,u,v,x,tot;
int c[maxn],fa[maxn],son[maxn],size[maxn],deep[maxn],rt[maxn],in[maxn],fin[maxn];
int lc[maxn*4],rc[maxn*4],sum[maxn*4],lazy[maxn*4];
vector<int> edge[maxn];
char s[5];
void dfs(int x,int pre)
{
    fa[x] = pre;
    son[x] = -1;
    size[x] = 1;
    deep[x] = deep[pre]+1;
    for(int i = 0; i < edge[x].size(); i++)
        if(edge[x][i] != pre)
        {
            dfs(edge[x][i],x);
            size[x] += size[edge[x][i]];
            if(son[x]==-1 || size[edge[x][i]]>size[son[x]]) son[x]=edge[x][i];
        }
}
void init(int x,int root)
{
    rt[x] = root;
    in[x] = ++tot;
    fin[in[x]] = x;
    if(son[x]==-1)return;
    init(son[x],root);
    for(int i = 0; i < edge[x].size(); i++)
        if(edge[x][i] != fa[x] && edge[x][i]!=son[x])
            init(edge[x][i],edge[x][i]);
}
void maintain(int o,int l,int r)
{
    if(l==r)return;
    lc[o]=lc[lson];
    rc[o]=rc[rson];
    sum[o]=sum[lson]+sum[rson]-(rc[lson]==lc[rson]);
}
void pushdown(int o,int l,int r)
{
    if(lazy[o]!=-1 && l!=r)
    {
        lc[lson]=lc[rson]=lazy[o];
        rc[lson]=rc[rson]=lazy[o];
        sum[lson]=sum[rson]=1;
        lazy[lson]=lazy[rson]=lazy[o];
    }
    lazy[o]=-1;
}
void buildtree(int o,int l,int r)
{
    if(l==r)
    {
        lc[o]=rc[o]=c[fin[l]];
        sum[o]=1;
        return;
    }
    int mid = (l+r)/2;
    buildtree(lson,l,mid);
    buildtree(rson,mid+1,r);
    maintain(o,l,r);
}
void addtree(int o,int l,int r,int L,int R,int c)
{
    pushdown(o,l,r);
    if(l>R || r<L)return;
    if(l>=L && r<=R)
    {
        lazy[o]=c;
        lc[o]=rc[o]=c;
        sum[o]=1;
        return;
    }
    int mid=(l+r)/2;
    addtree(lson,l,mid,L,R,c);
    addtree(rson,mid+1,r,L,R,c);
    maintain(o,l,r);
}
int getsum(int o,int l,int r,int L,int R)
{
    pushdown(o,l,r);
    if(l>R || r<L) return 0;
    if(l>=L && r<=R) return sum[o];
    int mid=(l+r)/2;
    int ret1=getsum(lson,l,mid,L,R);
    int ret2=getsum(rson,mid+1,r,L,R);
    maintain(o,l,r);
    if(ret1 && ret2)return ret1+ret2-(rc[lson]==lc[rson]);
    else return ret1+ret2;
}
int getcol(int o,int l,int r,int x)
{
    pushdown(o,l,r);
    if(l>x || r<x) return 0;
    if(l==r) return lc[o];
    int mid=(l+r)/2,ret=0;
    ret=getcol(lson,l,mid,x)+getcol(rson,mid+1,r,x);
    maintain(o,l,r);
    return ret;
}
void insert(int x,int y,int col)
{
    while(rt[x]!=rt[y])
    {
        if(deep[rt[x]]>=deep[rt[y]])addtree(1,1,tot,in[rt[x]],in[x],col),x=fa[rt[x]];
        else addtree(1,1,tot,in[rt[y]],in[y],col),y=fa[rt[y]];
    }
    if(deep[x]>=deep[y])addtree(1,1,tot,in[y],in[x],col);
    else addtree(1,1,tot,in[x],in[y],col);
}
int query(int x,int y)
{
    int ret=0,lxcol=-1,lycol=-1;
    while(rt[x]!=rt[y])
    {
        //cout<<in[rt[x]]<<' '<<in[x]<<' '<<in[rt[y]]<<' '<<in[y]<<endl;
        if(deep[rt[x]]>=deep[rt[y]])
            ret+=getsum(1,1,tot,in[rt[x]],in[x])-(getcol(1,1,tot,in[fa[rt[x]]])==getcol(1,1,tot,in[rt[x]])),x=fa[rt[x]];
        else ret+=getsum(1,1,tot,in[rt[y]],in[y])-(getcol(1,1,tot,in[fa[rt[y]]])==getcol(1,1,tot,in[rt[y]])),y=fa[rt[y]];
    }
    if(deep[x]>=deep[y])ret+=getsum(1,1,tot,in[y],in[x]);
    else ret+=getsum(1,1,tot,in[x],in[y]);
    return ret;
}
int main(void)
{
    memset(lazy,-1,sizeof(lazy));
    scanf("%d%d",&n,&m);
    for(int i = 1; i <= n; i++)
        scanf("%d",&c[i]);
    for(int i = 1; i < n; i++)
    {
        scanf("%d%d",&u,&v);
        edge[u].push_back(v);
        edge[v].push_back(u);
    }
    dfs(1,0);
    init(1,1);
    buildtree(1,1,tot);
    //cout<<"wtf"<<endl;
    for(int i = 1; i <= m; i++)
    {
        scanf("%s%d%d",s,&u,&v);
        if(s[0]=='C')
        {
            scanf("%d",&x);
            insert(u,v,x);
        }
        else printf("%d\n",query(u,v));
    }
}