bzoj 2243 2243: [SDOI2011]染色 树链剖分

直接树链剖分就可以啦。

#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
using namespace std;
const int MAXN = 110000;
typedef int ll;
int head[MAXN],to[MAXN * 2],nxt[MAXN * 2],q[MAXN],dep[MAXN],son[MAXN],siz[MAXN],top[MAXN],pos[MAXN],idx[MAXN],fa[MAXN],jl[MAXN];
int data[MAXN * 5],lcol[MAXN * 5],rcol[MAXN * 5],lzy[MAXN * 5],p[MAXN][32];
int n,m,a,b,c,sum,cnt;
int col[MAXN];
char s[10];
void add(int a,int b)
{
    nxt[++cnt] = head[a];
    to[cnt] = b;
    head[a] = cnt;
}
void updata(int k)
{
    if (rcol[k << 1] == lcol[k << 1 | 1]) 
        data[k] = data[k << 1] + data[k << 1 | 1] - 1;
    else 
        data[k] = data[k << 1] + data[k << 1 | 1];
    lcol[k] = lcol[k << 1];
    rcol[k] = rcol[k << 1 | 1];
}
void build(int k,int l,int r)
{
    lzy[k] = -1;
    if (l == r) 
    {
        data[k] = 1;
        lcol[k] = rcol[k] = col[idx[l]];
        return;
    }
    int midn = l + r >> 1;
    build(k << 1,l,midn);
    build(k << 1 | 1,midn + 1,r);
    updata(k);
}
void lca_init()
{
    for (int j = 1;(1 << j) <= n;j++)
        for (int i = 1;i <= n;i++)
            p[i][j] = p[p[i][j - 1]][j - 1];
}
void init()
{
    int que = 0,u,v;
    q[++que] = 1;
    dep[1] = 1;
    for (int i = q[que];i <= que;i++)
    {
        siz[u = q[i]] = 1;
        for (int o = head[u];o;o = nxt[o])
        {
            v = to[o];
            if (v == fa[u]) 
                continue;
            fa[v] = u;
            p[v][0] = u;
            dep[v] = dep[u] + 1;
            q[++que] = v;
        }
    }
    for (int i = que;i >= 1;i--)
    {
        u = q[i];
        siz[v = fa[u]] += siz[u];
        if (siz[u] > siz[son[v]]) 
            son[v] = u;
    }
    for (int i = 1;i <= que;i++)
    {
        if (top[u = q[i]]) 
            continue;
        for (v = u;v;v = son[v])
        {
            top[v] = u;
            idx[pos[v] = ++sum] = v;
        }
    }
    build(1,1,sum);
    lca_init();
}
int lca(int a,int b)
{
    if (dep[a] < dep[b]) swap(a,b);
    for (int j = log2(n);j >= 0;j--)
        if (dep[a] - (1 << j) >= dep[b])
            a = p[a][j];
    if (a == b) return b;
    for (int j = log2(n);j >= 0;j--)
        if (p[a][j] != p[b][j])
        {
            a = p[a][j];
            b = p[b][j];
        }
    return p[a][0];
}
void push_down(int k)
{
    data[k << 1] = data[k << 1 | 1] = 1;
    lzy[k << 1] = lzy[k << 1 | 1] = lzy[k];
    lcol[k << 1] = rcol[k << 1] = lcol[k << 1 | 1] = rcol[k << 1 | 1] = lzy[k];
    lzy[k] = -1;
}
void vec_chg(int k,int l,int r,int x,int y,int xg)
{
    if (x <= l && r <= y)
    {
        lzy[k] = xg;
        lcol[k] = rcol[k] = xg;
        data[k] = 1;
        return;
    }
    if (lzy[k] >= 0) push_down(k);
    int midn = l + r >> 1;
    if (x <= midn) vec_chg(k << 1,l,midn,x,y,xg);
    if (y > midn) vec_chg(k << 1 | 1,midn + 1,r,x,y,xg);
    updata(k);
}
int vec_qry(int k,int l,int r,int x,int y)
{
    if (x <= l && r <= y) return data[k];
    if (lzy[k] >= 0) push_down(k);
    int midn = l + r >> 1,res = 0;
    if (x <= midn) res += vec_qry(k << 1,l,midn,x,y);
    if (y > midn) res += vec_qry(k << 1 | 1,midn + 1,r,x,y);
    if (x <= midn && y > midn) if (rcol[k << 1] == lcol[k << 1 | 1]) res--;
    return res;
}
int vec_col(int k,int l,int r,int x,int y)
{
    if (x <= l && r <= y) return lcol[k];
    if (lzy[k] >= 0) push_down(k);
    int midn = l + r >> 1;
    if (x <= midn) return vec_col(k << 1,l,midn,x,y);
    if (y > midn) return vec_col(k << 1 | 1,midn + 1,r,x,y);
}
void query_chg(int a,int b,int xg)
{
    while (top[a] != top[b])
    {
        if (dep[top[a]] < dep[top[b]]) swap(a,b);
        vec_chg(1,1,sum,pos[top[a]],pos[a],xg);
        a = fa[top[a]];
    }
    if (dep[a] > dep[b]) swap(a,b);
    vec_chg(1,1,sum,pos[a],pos[b],xg);
}
int query_sum(int a,int b)
{
    int res = 0;
    while (top[a] != top[b])
    {
        if (dep[top[a]] < dep[top[b]]) swap(a,b);
        res += vec_qry(1,1,sum,pos[top[a]],pos[a]);
        if (vec_col(1,1,sum,pos[top[a]],pos[top[a]]) == vec_col(1,1,sum,pos[fa[top[a]]],pos[fa[top[a]]])) res--;
        a = fa[top[a]];
    }
    if (dep[a] > dep[b]) swap(a,b);
    res += vec_qry(1,1,sum,pos[a],pos[b]);
    return res;
}
int main()
{
    scanf("%d%d",&n,&m);
    for (int i = 1;i <= n;i++) 
        scanf("%d",&col[i]);
    for (int i = 1;i <= n - 1;i++) 
    {   
        scanf("%d%d",&a,&b);
        add(a,b);
        add(b,a);
    }
    init();
    for (int i = 1;i <= m;i++)
    {
        scanf("%s",s);
        if (s[0] == 'C')
        {
            scanf("%d%d%d",&a,&b,&c);
            query_chg(a,b,c);
        }else
        {
            scanf("%d%d",&a,&b);
            int tt = lca(a,b);
            printf("%d\n",query_sum(a,tt) + query_sum(b,tt) - 1);
        }
    }
    return 0;
}