luogu2486 [SDOI2011]染色

link

题目大意：给定一个N个点的树，每个点有一个颜色
有M次操作，每次可以修改树某条链所有点变成一个颜色，查询某条链上点的颜色段数

树剖，线段树维护区间合并

我的代码记录的是某个区间左端点颜色、右端点颜色、除了左端点和右端点的颜色段数

需要稍微特殊处理一些情况，详见代码

#include <cstdio>
#include <vector>
using namespace std;

struct fuck
{
    int l, r, cnt; //cnt=-1代表lcol=rcol 区间内只有一种颜色，-2代表区间是空的
    fuck(){}
    fuck(int col) : l(col), r(col), cnt(-1) {}
    fuck(int l, int r, int cnt) : l(l), r(r), cnt(cnt) {}
};

fuck operator*(const fuck &l, const fuck &r)
{
    if (l.cnt == -2) return r;
    if (r.cnt == -2) return l;
    if (l.cnt == -1 && r.cnt == -1)
    {
        if (l.l == r.l) return fuck(l.l, l.l, -1);
        else return fuck(l.l, r.l, 0);
    }
    else if (l.cnt == -1)
    {
        if (l.l == r.l) return r;
        else return fuck(l.l, r.r, r.cnt + 1);
    }
    else if (r.cnt == -1)
    {
        if (l.r == r.l) return l;
        else return fuck(l.l, r.l, l.cnt + 1);
    }
    else
    {
        return fuck(l.l, r.r, l.cnt + r.cnt + 2 - (l.r == r.l));
    }
}

vector<int> out[100010];
int n, m, col[100010];
int fa[100010], depth[100010], weight[100010], wson[100010];
int dfn[100010], top[100010], id[100010], tot;
fuck tree[400010]; int lazy[400010];

fuck rev(const fuck &x)
{
    return fuck(x.r, x.l, x.cnt);
}

void dfs1(int x)
{
    weight[x] = 1, wson[x] = -1;
    for (int i : out[x])
        if (fa[x] != i)
        {
            fa[i] = x, depth[i] = depth[x] + 1;
            dfs1(i);
            weight[x] += weight[i];
            if (wson[x] == -1 || weight[i] > weight[wson[x]]) wson[x] = i;
        }
}

void dfs2(int x, int topf)
{
    dfn[x] = ++tot, top[x] = topf, id[tot] = x;
    if (wson[x] != -1)
    {
        dfs2(wson[x], topf);
        for (int i : out[x])
            if (i != fa[x] && i != wson[x])
                dfs2(i, i);
    }
}

void build(int x, int cl, int cr)
{
    if (cl == cr)
    {
        tree[x] = fuck(col[id[cl]]);
        return;
    }
    int mid = (cl + cr) / 2;
    build(x * 2, cl, mid);
    build(x * 2 + 1, mid + 1, cr);
    tree[x] = tree[x * 2] * tree[x * 2 + 1];
}

void pushdown(int x)
{
    if (lazy[x])
    {
        tree[x * 2] = tree[x * 2 + 1] = fuck(lazy[x]);
        lazy[x * 2] = lazy[x * 2 + 1] = lazy[x];
        lazy[x] = 0;
    }
}

void chenge(int x, int cl, int cr, int L, int R, int col)
{
    if (R < cl || cr < L) return;
    if (L <= cl && cr <= R)
    {
        tree[x] = fuck(col);
        lazy[x] = col;
        return;
    }
    pushdown(x);
    int mid = (cl + cr) / 2;
    chenge(x * 2, cl, mid, L, R, col);
    chenge(x * 2 + 1, mid + 1, cr, L, R, col);
    tree[x] = tree[x * 2] * tree[x * 2 + 1];
}

fuck query(int x, int cl, int cr, int L, int R)
{
    if (R < cl || cr < L) return fuck(0, 0, -2);
    if (L <= cl && cr <= R) return tree[x];
    pushdown(x);
    int mid = (cl + cr) / 2;
    return query(x * 2, cl, mid, L, R) * query(x * 2 + 1, mid + 1, cr, L, R);
}

int main()
{
    freopen("paint.in", "r", stdin);
    freopen("paint.out", "w", stdout);
    scanf("%d%d", &n, &m);
    for (int i = 1; i <= n; i++) scanf("%d", &col[i]);
    for (int x, y, i = 1; i < n; i++) scanf("%d%d", &x, &y), out[x].push_back(y), out[y].push_back(x);
    dfs1(1), dfs2(1, 1);
    build(1, 1, n);
    char ch;
    for (int x, y, z, i = 1; i <= m; i++)
    {
        scanf(" %c%d%d", &ch, &x, &y);
        if (ch == 'C')
        {
            scanf("%d", &z);
            while (top[x] != top[y])
            {
                if (depth[top[x]] < depth[top[y]]) swap(x, y);
                chenge(1, 1, n, dfn[top[x]], dfn[x], z);
                x = fa[top[x]];
            }
            if (depth[x] > depth[y]) swap(x, y);
            chenge(1, 1, n, dfn[x], dfn[y], z);
        }
        else
        {
            fuck ans1(0, 0, -2), ans2(0, 0, -2);
            while (top[x] != top[y])
            {
                if (depth[top[x]] < depth[top[y]]) swap(x, y), swap(ans1, ans2);
                ans1 = query(1, 1, n, dfn[top[x]], dfn[x]) * ans1;
                x = fa[top[x]];
            }
            if (depth[x] > depth[y]) swap(x, y), swap(ans1, ans2);
            fuck ans = rev(ans1) * query(1, 1, n, dfn[x], dfn[y]) * ans2;
            printf("%d\n", ans.cnt + 2);
        }
    }
    fclose(stdin);
    fclose(stdout);
    return 0;
}