洛谷 P2486 [SDOI2011]染色 树链剖分

题目：
洛谷 P2486

大意：一棵树，每个点有颜色，支持两种操作，链上修改，链上查询颜色段。

分析：树链剖分一下，线段树打个标记就可以了。

代码：

#include <iostream>
#include <cstdio>
#include <cmath>

const int maxn=100005;

using namespace std;

int test,n,cnt,size[maxn],top[maxn],fa[maxn],dep[maxn],last[maxn],a[maxn],sz,pos[maxn];
struct data{int y,next;}edge[maxn*2];
struct tree{int l,r,s,c;}t[maxn*4];

void insert_edge(int u,int v)
{
    edge[++cnt].y=v;edge[cnt].next=last[u];last[u]=cnt;
    edge[++cnt].y=u;edge[cnt].next=last[v];last[v]=cnt;
}

void initt()
{
    int i,u,v;
    scanf("%d%d",&n,&test);
    for (i=1;i<=n;i++)
        scanf("%d",&a[i]);
    for (i=1;i<=n-1;i++)
    {
        scanf("%d%d",&u,&v);
        insert_edge(u,v);
    }
}

void dfs1(int x)
{
    int c=last[x];
    size[x]=1;
    while (c>0)
    {
        int y=edge[c].y;
        if (y!=fa[x])
        {
            fa[y]=x;
            dep[y]=dep[x]+1;
            dfs1(y);
            size[x]+=size[y];
        }
        c=edge[c].next;

    }
}

void dfs2(int x,int ch)
{
    int k=0; sz++;
    pos[x]=sz;
    top[x]=ch;
    for (int c=last[x];c>0;c=edge[c].next)
    {
        int y=edge[c].y;
        if ((fa[y]==x) && (size[y]>size[k]))
        {
            k=y;
        } 
    }
    if (k==0) return;
    dfs2(k,ch);
    for (int c=last[x];c>0;c=edge[c].next)
    {
        int y=edge[c].y;
        if ((fa[y]==x) && (y!=k)) dfs2(y,y);
    }
}

void change(int p,int l,int r,int x,int y,int k)
{
    if ((l==x) && (r==y))
    {
        t[p].l=k;
        t[p].r=k;
        t[p].s=1;
        t[p].c=1;
        return;
    }
    if (t[p].c==1)
    {
        t[2*p].c=1;
        t[2*p].l=t[p].l;
        t[2*p].r=t[p].r;
        t[2*p].s=1;
        t[2*p+1]=t[2*p];
        t[p].c=0;
    }
    int mid=(l+r)/2;
    if (y<=mid) change(p*2,l,mid,x,y,k);
    else
    {
        if (x>mid) change(p*2+1,mid+1,r,x,y,k);
        else
        {
            change(p*2,l,mid,x,mid,k);
            change(p*2+1,mid+1,r,mid+1,y,k);
        }
    }
    t[p].s=t[p*2].s+t[p*2+1].s;
    if (t[p*2].r==t[p*2+1].l) t[p].s--;
    t[p].l=t[p*2].l;
    t[p].r=t[p*2+1].r;
}

void getsum(int p,int l,int r,int x,int y,int &lc,int &rc,int &sum)
{
    if ((l==x) && (r==y))
    {
        lc=t[p].l;
        rc=t[p].r;
        sum=t[p].s;
        return;
    }
    if (t[p].c==1)
    {
        lc=t[p].l; rc=t[p].r;
        sum=1;
        return;
    }
    int mid=(l+r)/2;
    if (y<=mid) getsum(p*2,l,mid,x,y,lc,rc,sum);
    else
    {
        if (x>mid) getsum(p*2+1,mid+1,r,x,y,lc,rc,sum);
        else
        {
            getsum(p*2,l,mid,x,mid,lc,rc,sum);
            int re_lc=lc,re_rc=rc,re_sum=sum;
            getsum(p*2+1,mid+1,r,mid+1,y,lc,rc,sum);
            sum+=re_sum;
            if (re_rc==lc) sum--;
            lc=re_lc;
        }
    }
}

void change_line(int x,int y,int k)
{
    while (top[x]!=top[y])
    {
        if (dep[top[x]]>dep[top[y]]) swap(x,y);
        change(1,1,n,pos[top[y]],pos[y],k);
        y=fa[top[y]];
    }
    if (dep[x]>dep[y]) swap(x,y);
    change(1,1,n,pos[x],pos[y],k);
}

int ask(int x,int y)
{   
    int ans=0;
    int colx=-1,coly=-1;
    int lc,rc,t;    
    while (top[x]!=top[y])
    {
        if (dep[top[x]]>dep[top[y]]) 
        {
            swap(x,y);
            swap(colx,coly);
        }
        lc=0,rc=0,t=0;
        getsum(1,1,n,pos[top[y]],pos[y],lc,rc,t);
        if (rc==coly) ans+=t-1;
                 else ans+=t;
        coly=lc;
        y=fa[top[y]];
    }
    if (dep[x]>dep[y]) {swap(x,y); swap(colx,coly);}

    getsum(1,1,n,pos[x],pos[y],lc,rc,t);
    ans+=t;
    if (lc==colx) ans--;
    if (rc==coly) ans--;
    return ans;
}

int main()
{
    initt();
    dfs1(1);
    dfs2(1,1);   
    for (int i=1;i<=n;i++)
    {
        change(1,1,n,pos[i],pos[i],a[i]);
    }  
    char ch[10]; int x,y,z;      
    for (int i=1;i<=test;i++)
    {       
        scanf("%s",ch);
        if (ch[0]=='C')
        {
            scanf("%d%d%d",&x,&y,&z);
            change_line(x,y,z); 
        }
        else
        {
            scanf("%d%d",&x,&y);
            printf("%d\n",ask(x,y));
        }                      
    }
    return 0;
}