[BZOJ2243][SDOI2011]染色

(fx，x写反了，还有个地方没加括号，都能过样例，RE好多发)

比较简单的树剖，需要注意的地方是因为分成了几段区间求解，所以每次都需要判断边界颜色是否相同，相同答案就要减1。

代码还是有点长，跑得还算快

/**************************************************************
    Problem: 2243
    User: cabinfever
    Language: C++
    Result: Accepted
    Time:4560 ms
    Memory:34412 kb
****************************************************************/
 
#include <cstdio>
#include <cstring>
#include <string>
#include <cstdlib>
#include <cmath>
#include <iostream>
#include <algorithm>
 
using namespace std;
 
const int maxn = 300010;
 
struct edge
{
    int v,next;
}e[maxn << 1];
int h[maxn],num = 0;
struct node
{
    int l,r,mark,x;
    int rc,lc;
}t[maxn << 1];
int tot = 0;
int a[maxn];
bool vis[maxn];
int n,m,u,v;
int cnt = 0;
int fa[maxn];
int dep[maxn];
int son[maxn];
int pos[maxn],ppos[maxn];
int poi[maxn];
int top[maxn];
int size[maxn];
int ll,rr;
 
void build_edge(int u,int v)
{
    num++;
    e[num].v = v;
    e[num].next = h[u];
    h[u] = num;
}
 
void dfs1(int x)
{
    dep[x] = dep[fa[x]] + 1;
    vis[x] = true;
    size[x] = 1;
    son[x] = 0;
    for(int i = h[x]; i; i = e[i].next)
    {
        if(vis[e[i].v])
            continue;
        fa[e[i].v] = x;
        dfs1(e[i].v);
        if(size[e[i].v] > size[son[x]])
            son[x] = e[i].v;
        size[x] += size[e[i].v];
    }
}
 
void dfs2(int x)
{
    vis[x] = true;
    if(son[fa[x]] == x)
        top[x] = top[fa[x]];
    else
    {
        top[x] = x;
        for(int i = x; i; i = son[i])
        {
            pos[i] = ++cnt;
            ppos[cnt] = i;
        }
    }
    for(int i = h[x]; i; i = e[i].next)
        if(!vis[e[i].v])
            dfs2(e[i].v);
}
 
void build(int p,int l,int r)
{
    int mid = l + r >> 1;
    if(l == r)
    {
        t[p].x = 1;
        t[p].lc = a[ppos[mid]];
        t[p].rc = a[ppos[mid]];
        poi[mid] = p;
        return;
    }
    t[p].l = ++tot;
    t[p].r = ++tot;
    build(t[p].l,l,mid);
    build(t[p].r,mid+1,r);
    t[p].x = t[t[p].l].x + t[t[p].r].x;
    if(t[t[p].l].rc == t[t[p].r].lc)
        t[p].x--;
    t[p].lc = t[t[p].l].lc;
    t[p].rc = t[t[p].r].rc;
}
 
void change(int p,int l,int r,int L,int R,int c)
{
    int mid = l + r >> 1;
    if(l == L && r == R)
    {
        t[p].x = 1;
        t[p].rc = c;
        t[p].lc = c;
        t[p].mark = c;
        return;
    }
    if(t[p].mark)
    {
        if(t[p].mark == c)
            return;
        t[t[p].l].x = 1;
        t[t[p].l].rc = t[p].mark;
        t[t[p].l].lc = t[p].mark;
        t[t[p].l].mark = t[p].mark;
        t[t[p].r].x = 1;
        t[t[p].r].rc = t[p].mark;
        t[t[p].r].lc = t[p].mark;
        t[t[p].r].mark = t[p].mark;
        t[p].mark = 0;
    }
    if(R <= mid)
        change(t[p].l,l,mid,L,R,c);
    else if(L > mid)
        change(t[p].r,mid+1,r,L,R,c);
    else
    {
        change(t[p].l,l,mid,L,mid,c);
        change(t[p].r,mid+1,r,mid+1,R,c);
    }
    t[p].x = t[t[p].l].x + t[t[p].r].x;
    if(t[t[p].l].rc == t[t[p].r].lc)
        t[p].x--;
    t[p].lc = t[t[p].l].lc;
    t[p].rc = t[t[p].r].rc;
}
 
int getans(int p,int l,int r,int L,int R)
{
    int mid = l + r >> 1;
    if(l == L && r == R)
    {
        if(r == u)
            rr = t[p].rc;
        if(l == v)
            ll = t[p].lc;
        return t[p].x;
    }
    if(t[p].mark)
    {
        t[t[p].l].x = 1;
        t[t[p].l].rc = t[p].mark;
        t[t[p].l].lc = t[p].mark;
        t[t[p].l].mark = t[p].mark;
        t[t[p].r].x = 1;
        t[t[p].r].rc = t[p].mark;
        t[t[p].r].lc = t[p].mark;
        t[t[p].r].mark = t[p].mark;
        t[p].mark = 0;
    }
    if(R <= mid)
        return getans(t[p].l,l,mid,L,R);
    else if(L > mid)
        return getans(t[p].r,mid+1,r,L,R);
    int ans = getans(t[p].l,l,mid,L,mid) + getans(t[p].r,mid+1,r,mid+1,R);
    if(t[t[p].l].rc == t[t[p].r].lc)
        ans--;
    return ans;
}
 
int query(int x,int y)
{
    int ans = 0;
    int fx = top[x];
    int fy = top[y];
    int xc = -1, yc = -1;
    while(fx != fy)
    {
        if(dep[fx] > dep[fy])
        {
            u = pos[x];
            v = pos[fx];
            ans += getans(0,1,n,pos[fx],pos[x]);
            if(rr == xc)
                ans--;
            xc = ll;
            x = fa[fx];
            fx = top[x];
        }
        else
        {
            u = pos[y];
            v = pos[fy];
            ans += getans(0,1,n,pos[fy],pos[y]);
            if(rr == yc)
                ans--;
            yc = ll;
            y = fa[fy];
            fy = top[y];
        }
    }
    if(pos[x] > pos[y])
        u = pos[x], v = pos[y], ans += getans(0,1,n,pos[y],pos[x]);
    else
        u = pos[y], v = pos[x], ans += getans(0,1,n,pos[x],pos[y]);
    if(pos[x] > pos[y])
        swap(ll,rr);
    if(ll == xc)
        ans--;
    if(rr == yc)
        ans--;
    return ans;
}
 
void updata(int x,int y,int c)
{
    int fx = top[x];
    int fy = top[y];
    while(fx != fy)
    {
        if(dep[fx] > dep[fy])
        {
            change(0,1,n,pos[fx],pos[x],c);
            x = fa[fx];
            fx = top[x];
        }
        else
        {
            change(0,1,n,pos[fy],pos[y],c);
            y = fa[fy];
            fy = top[y];
        }
    }
    if(pos[x] > pos[y])
        change(0,1,n,pos[y],pos[x],c);
    else
        change(0,1,n,pos[x],pos[y],c);
}
 
int main()
{
    scanf("%d%d",&n,&m);
    for(int i = 1; i <= n; i++)
        scanf("%d",&a[i]);
    for(int i = 0; i < n; i++)
    {
        scanf("%d%d",&u,&v);
        build_edge(u,v);
        build_edge(v,u);
    }
    memset(vis,0,sizeof(vis));
    dfs1(1);
    memset(vis,0,sizeof(vis));
    dfs2(1);
    build(0,1,n);
    char q[10];
    int c;
    while(m--)
    {
        scanf("%s",q);
        if(q[0] == 'Q')
        {
            scanf("%d%d",&u,&v);
            printf("%d\n",query(u,v));
        }
        else
        {
            scanf("%d%d%d",&u,&v,&c);
            updata(u,v,c);
        }
    }
    return 0;
}