bzoj2243: [SDOI2011]染色 树链剖分

熟练度很重要，一次AC~~~。

/**************************************************************
    Problem: 2243
    User: xujiahe
    Language: C++
    Result: Accepted
    Time:6748 ms
    Memory:20908 kb
****************************************************************/
 
#include <iostream>
#include <cstring>
#include <cstdio>
#include <algorithm>
#include <cmath>
#include <queue>
using namespace std;
#define maxn 200000
#define ll long long
#define lson rt<<1,l,m
#define rson rt<<1|1,m+1,r
#define  ss printf("orz\n")
int fa[maxn],siz[maxn],top[maxn],w[maxn],son[maxn],dep[maxn],id,lc,rc;
struct node
{
    int u,v,next;
}g[maxn];
int num,n,first[maxn],val[maxn];
int tree[maxn<<1],col[maxn<<1],lcol[maxn<<1],rcol[maxn<<1];
void build(int a,int b)
{
    num++;
    g[num].v=b;
    g[num].next=first[a];
    first[a]=num;
}
void dfs(int now)
{
    int v,ceng;
    int maxv=0,flag=0;
    siz[now]=1;
    for(int i=first[now];i;i=g[i].next)
    {
        v=g[i].v;
        if(v==fa[now]) continue;
        dep[v]=dep[now]+1;
        fa[v]=now;
        dfs(v);
        siz[now]+=siz[v];
        if(siz[v]>maxv)
        {
            maxv=siz[v];
            flag=v;
        }
    }
    son[now]=flag;
}
void getid(int now,int root)
{
    w[now]=++id;
    top[now]=root;
    if(son[now]) getid(son[now],top[now]);
    for(int i=first[now]; i; i=g[i].next)
    {
        if(g[i].v!=son[now]&&g[i].v!=fa[now])
        {
            getid(g[i].v,g[i].v);
        }
    }
}
void print(int rt,int l,int r)
{
    int m=(l+r)>>1;
    printf("%d %d %d\n",rt,l,r);
    printf("%d\n",col[rt]);
    printf("%d %d %d\n\n",tree[rt],lcol[rt],rcol[rt]);
    if(l==r)
    {
        return ;
    }
    print(lson);
    print(rson);
}
void pushdown(int rt)
{
    if(col[rt]!=-1)
    {
        col[rt<<1]=col[rt<<1|1]=col[rt];
        lcol[rt<<1]=rcol[rt<<1]=lcol[rt<<1|1]=rcol[rt<<1|1]=col[rt];
        tree[rt<<1]=tree[rt<<1|1]=1;
        col[rt]=-1;
    }
}
void pushup(int rt)
{
    tree[rt]=tree[rt<<1]+tree[rt<<1|1];
    lcol[rt]=lcol[rt<<1];
    rcol[rt]=rcol[rt<<1|1];
    if(rcol[rt<<1]==lcol[rt<<1|1])
        tree[rt]--;
}
void update(int rt,int l,int r,int x,int y,int add)
{
    if(x<=l&&y>=r)
    {
        tree[rt]=1;
        lcol[rt]=rcol[rt]=add;
        col[rt]=add;
        return ;
    }
    pushdown(rt);
    int m=(l+r)>>1;
    if(x<=m)
    {
        update(lson,x,y,add);
    }
    if(y>m)
    {
        update(rson,x,y,add);
    }
    pushup(rt);
}
int query(int rt,int l,int r,int x,int y)
{
    if(l==x)
    {
        lc=lcol[rt];
    }
    if(r==y)
    {
        rc=rcol[rt];
    }
    if(x<=l&&y>=r)
    {
        return tree[rt];
    }
    int tmp=0;
    int t1,t2;
    t1=t2=-1;
    pushdown(rt);
    int m=(l+r)>>1;
    if(x<=m)
    {
        t1=rcol[rt<<1];
        tmp+=query(lson,x,y);
    }
    if(y>m)
    {
        t2=lcol[rt<<1|1];
        tmp+=query(rson,x,y);
    }
    if(t1==t2&&t1!=-1)
    {
        tmp--;
    }
    return tmp;
}
void getans(int x,int y,int z,int op)
{
    int f1,f2,t1,t2,t;
    t1=t2=-1;
    int ans=0;
    if(op==1)
    {
        while(top[x]!=top[y])
        {
            if(dep[top[x]]<dep[top[y]])
            {
                swap(x,y);
                swap(t1,t2);
            }
            ans+=query(1,1,n,w[top[x]],w[x]);
            ans-= (t1== rc);
            t1 = lc;
            x=fa[top[x]];
        }
        if (dep[x] < dep[y])
        {
            swap(x,y);
            swap(t1,t2);
        }
        ans+=query(1,1,n,w[y],w[x]);
        ans -= ((t1 == rc && t1 != -1) + (t2 == lc && t2 != -1));
        printf("%d\n",ans);
        return ;
    }
    if(op==2)
    {
 
        while(top[x]!=top[y])
        {
            if(dep[top[x]]<dep[top[y]])
            {
                swap(x,y);
            }
            update(1,1,n,w[top[x]],w[x],z);
            x=fa[top[x]];
        }
        if (dep[x] < dep[y])
        {
            swap(x,y);
        }
        update(1,1,n,w[y],w[x],z);
        return ;
    }
}
int main()
{
    char str[20];
    int a,b,m;
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&val[i]);
    }
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&a,&b);
        build(a,b);
        build(b,a);
    }
    dfs(1);
    getid(1,1);
    int x,y,z;
    memset(col,-1,sizeof(col));
    for(int i=1;i<=n;i++)
    {
        update(1,1,n,w[i],w[i],val[i]);
    }
    while(m--)
    {
        scanf("%s",str);
        if(str[0]=='Q')
        {
            scanf("%d%d",&x,&y);
            getans(x,y,0,1);
        }
        else
        {
            scanf("%d%d%d",&x,&y,&z);
            getans(x,y,z,2);
        }
    }
    return 0;
}
/*
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
*/