BZOJ2243: [SDOI2011]染色

题目链接
树链剖分，这题两个区间合并时和Query求值时都需要注意！！！

#include<iostream>
#include<cstdio>
using namespace std;
int n,m,tot,lable,color[410000],Next[410000],head[410000],tree[410000],fa[410000],dep[410000],size[410000];
int son[410000],tid[410000],number[410000],top[410000],a[880000],L[880000],R[880000],lazy[880000];
bool visit[410000];
void add(int x,int y)
{
    tot++;
    Next[tot]=head[x];
    head[x]=tot;
    tree[tot]=y;
}
void dfs(int x,int depth,int father)
{
    visit[x]=true;fa[x]=father;dep[x]=depth;size[x]=1;son[x]=0;
    int maxsize=0;
    for (int i=head[x];i;i=Next[i])
    if (!visit[tree[i]])
    {
        dfs(tree[i],depth+1,x);
        size[x]+=size[tree[i]];
        if (size[tree[i]]>maxsize)
        {
            maxsize=size[tree[i]];
            son[x]=tree[i];
        }
    }
}
void dfs1(int x,int ancestor)
{
    visit[x]=true;tid[x]=++lable;number[lable]=x;top[x]=ancestor;
    if (son[x]!=0) dfs1(son[x],ancestor);
    for (int i=head[x];i;i=Next[i]) 
        if (!visit[tree[i]]) dfs1(tree[i],tree[i]);
}
void up(int id)
{
    a[id]=a[id*2]+a[id*2+1];
    if (R[id*2]==L[id*2+1]) a[id]--;
    L[id]=L[id*2];R[id]=R[id*2+1];
}
void down(int id)
{
    if (lazy[id]!=-1)
    {
        L[id*2]=R[id*2]=L[id*2+1]=R[id*2+1]=lazy[id];
        lazy[id*2]=lazy[id*2+1]=lazy[id];
        lazy[id]=-1;
        a[id*2]=a[id*2+1]=1;
    }
}
void swap(int &x,int &y)
{
    int t=x;x=y;y=t;
}
void build(int id,int l,int r)
{
    lazy[id]=-1;L[id]=-1;R[id]=-1;a[id]=0;
    if (l==r) 
    {
        a[id]=1;
        L[id]=R[id]=color[number[l]];
        return;
    }
    int mid=(l+r)/2;
    build(id*2,l,mid);
    build(id*2+1,mid+1,r);
    up(id);
}
int query(int x,int y,int id,int l,int r)
{
    if (l>y||r<x) return -1;
    down(id);
    if (x<=l&&r<=y) return a[id];
    int mid=(l+r)/2;
    int x1=query(x,y,id*2,l,mid);
    int y1=query(x,y,id*2+1,mid+1,r);
    if (x1==-1) return y1;
    if (y1==-1) return x1;
    if (R[id*2]==L[id*2+1]) return x1+y1-1;
    return x1+y1;
}
int find(int id,int l,int r,int x,int y)
{
    down(id);
    if (l==r)
    {
        if (y==0) return L[id];else return R[id];
    }
    int mid=(l+r)/2;
    if (x<=mid) return find(id*2,l,mid,x,y);
        else return find(id*2+1,mid+1,r,x,y);
}
int Query(int x,int y)
{
    int ans=0;
    while (top[x]!=top[y])
    {
        if (dep[top[x]]<dep[top[y]]) swap(x,y);
        ans+=query(tid[top[x]],tid[x],1,1,n);
        if (find(1,1,n,tid[top[x]],0)==find(1,1,n,tid[fa[top[x]]],1)) ans--;
        x=fa[top[x]];
    }
    if (dep[x]>dep[y]) swap(x,y);
    ans+=query(tid[x],tid[y],1,1,n);
    return ans;
}
void change(int x,int y,int id,int l,int r,int d)
{
    if (l>y||r<x) return;
    down(id);
    if (x<=l&&r<=y) 
    {
        a[id]=1;
        L[id]=R[id]=lazy[id]=d;
        return;
    }
    int mid=(l+r)/2;
    change(x,y,id*2,l,mid,d);
    change(x,y,id*2+1,mid+1,r,d);
    up(id);
}
void Change(int x,int y,int d)
{
    while (top[x]!=top[y])
    {
        if (dep[top[x]]<dep[top[y]]) swap(x,y);
        change(tid[top[x]],tid[x],1,1,n,d);
        x=fa[top[x]];
    }
    if (dep[x]>dep[y]) swap(x,y);
    change(tid[x],tid[y],1,1,n,d);
}
int main()
{
    scanf("%d%d",&n,&m);
    tot=lable=0;
    for (int i=1;i<=n;i++) scanf("%d",&color[i]);
    for (int i=1;i<n;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);add(y,x);
    }
    for (int i=1;i<=n;i++) visit[i]=false;
    dfs(1,1,0);
    for (int i=1;i<=n;i++) visit[i]=false;
    dfs1(1,1);
    build(1,1,n);
    char s[100];
    for (int i=1;i<=m;i++)
    {
        scanf("%s",s);
        if (s[0]=='Q')
        {
            int x,y;
            scanf("%d%d",&x,&y);
            printf("%d\n",Query(x,y));
        }else
        {
            int x,y,z;
            scanf("%d%d%d",&x,&y,&z);
            Change(x,y,z);
        }
    }
    return 0;
}