[题解]luogu P2486 [SDOI2011]染色 树剖

戳

对于树剖里的线段树，我们多维护一个当前区间左端点的颜色和当前区间右端点的颜色
所以\(pushup\)时就变成了这样:

t[p].lc = t[p << 1].lc, t[p].rc = t[p << 1 | 1].rc;
t[p].s = t[p << 1].s + t[p << 1 | 1].s;
if (t[p << 1].rc == t[p << 1 | 1].lc) --t[p].s;

还要注意注意在我们跳链时，假如当前id[top[u]]的颜色与id[fa[top[u]]]的颜色相同，我们要记得给答案减1，因为它们是同种颜色属于同一段而我们把这段切成两段加了两次

#include <cstdio>
#include <algorithm>

const int N = 1e5 + 30;

int n, m, len, cnt, x, a, b, c;
int w[N], f[N], dep[N], fa[N], top[N], siz[N], son[N], dfn[N], id[N];
char ch[5];

struct Edge
{
    int to, nx;
}e[N << 1];
struct SegmentTree
{
    int l, r, s, lc, rc, lazy;
}t[N << 2];

int read()
{
    int s = 0, ff = 1;
    char cc = getchar();
    while (cc < '0' || cc > '9')
    {
        if (cc == '-') ff = -1;
        cc = getchar();
    }
    while (cc >= '0' && cc <= '9')
    {
        s = s * 10 + cc - '0';
        cc = getchar();
    }
    return s * ff;
}

inline void add_edge(int from, int to)
{
    e[++len] = (Edge){to, f[from]};
    f[from] = len;
}

void dfs1(int p, int from)
{
    fa[p] = from, dep[p] = dep[fa[p]] + 1, siz[p] = 1;
    for (int i = f[p]; i; i = e[i].nx)
    {   
        if (fa[p] == e[i].to) continue;
        dfs1(e[i].to, p), siz[p] += siz[e[i].to];
        if (siz[e[i].to] > siz[son[p]]) son[p] = e[i].to;
    }
}

void dfs2(int p, int tp)
{
    dfn[++cnt] = p, id[p] = cnt, top[p] = tp;
    if (son[p]) dfs2(son[p], tp);
    for (int i = f[p]; i; i = e[i].nx)
        if (e[i].to != son[p] && e[i].to != fa[p]) dfs2(e[i].to, e[i].to);
}

void build(int p, int l, int r)
{
    t[p].l = l, t[p].r = r;
    if (l == r) { t[p].s = 1; t[p].lc = t[p].rc = w[dfn[l]]; return ; }
    int mid = l + r >> 1;
    build(p << 1, l, mid), build(p << 1 | 1, mid + 1, r);
    t[p].lc = t[p << 1].lc, t[p].rc = t[p << 1 | 1].rc, t[p].s = t[p << 1].s + t[p << 1 | 1].s;
    if (t[p << 1].rc == t[p << 1 | 1].lc) --t[p].s;
}

inline void spread(int p)
{
    if (!t[p].lazy) return ;
    t[p << 1].s = t[p << 1 | 1].s = 1;
    t[p << 1].lazy = t[p << 1].lc = t[p << 1].rc = t[p].lazy;
    t[p << 1 | 1].lazy = t[p << 1 | 1].lc = t[p << 1 | 1].rc = t[p].lazy;
    t[p].lazy = 0;
}

void change(int p, int l, int r, int col)
{
    if (t[p].l >= l && t[p].r <= r) 
        { t[p].lc = t[p].rc = t[p].lazy = col, t[p].s = 1; return ; }
    spread(p);
    int mid = t[p].l + t[p].r >> 1;
    if (l <= mid) change(p << 1, l, r, col);
    if (mid < r) change(p << 1 | 1, l, r, col);
    t[p].lc = t[p << 1].lc, t[p].rc = t[p << 1 | 1].rc, t[p].s = t[p << 1].s + t[p << 1 | 1].s;
    if (t[p << 1].rc == t[p << 1 | 1].lc) --t[p].s;
}

int ask(int p, int l, int r)
{
    if (t[p].l >= l && t[p].r <= r) return t[p].s;
    spread(p);
    int mid = t[p].l + t[p].r >> 1, tot = 0;
    if (l <= mid && mid < r && t[p << 1].rc == t[p << 1 | 1].lc) --tot;
    if (l <= mid) tot += ask(p << 1, l, r);
    if (mid < r) tot += ask(p << 1 | 1, l, r);
    return tot;
}

int get(int p, int pos)
{
    if (t[p].l == t[p].r) return t[p].lc;
    spread(p);
    int mid = t[p].l + t[p].r >> 1;
    if (pos <= mid) return get(p << 1, pos);
    else return get(p << 1 | 1, pos);
}

inline void treechange(int u, int v, int col)
{
    while (top[u] != top[v])
    {
        if (dep[top[u]] < dep[top[v]]) std::swap(u, v);
        change(1, id[top[u]], id[u], col);
        u = fa[top[u]];
    }
    if (dep[u] < dep[v]) std::swap(u, v);
    change(1, id[v], id[u], col);
}

inline int treeask(int u, int v)
{
    int ans = 0, cu, cv;
    while (top[u] != top[v])
    {
        if (dep[top[u]] < dep[top[v]]) std::swap(u, v);
        ans += ask(1, id[top[u]], id[u]);
        cu = get(1, id[top[u]]), cv = get(1, id[fa[top[u]]]);
        if (cu == cv) --ans;
        u = fa[top[u]];
    }
    if (dep[u] < dep[v]) std::swap(u, v);
    ans += ask(1, id[v], id[u]);
    if (ans == 0) return 1;
    return ans;
}

int main()
{
    n = read(), m = read();
    for (int i = 1; i <= n; ++i) w[i] = read();
    for (int i = 1; i < n; ++i)
    {
        a = read(), b = read();
        add_edge(a, b), add_edge(b, a);
    }
    dfs1(1, 0), dfs2(1, 1), build(1, 1, n);
    for (int i = 1; i <= m; ++i)
    {
        scanf ("%s", ch);
        if (ch[0] == 'C')
        {
            a = read(), b = read(), c = read();
            treechange(a, b, c);
        }
        else
        {
            a = read(), b = read();
            printf ("%d\n", treeask(a, b));
        }    
    }
    return 0;
}