P2486 [SDOI2011]染色 维护区间块数 树链剖分

https://www.luogu.org/problemnew/show/P2486

题意

在树上维护两种操作，一种是把$x$到$y$间的点都染成$c$色，另一种是求$x$到$y$间的点有多少个颜色块，比如112221由“11”，“222”，“1”三块组成。

思路

这题的关键是要如何维护这个颜色块，我们可以利用线段树，记录每个区间的块数，区间左端点，区间右端点的颜色。合并中如果两个区间相邻点颜色相同，个数要减去1.
查询也是一样的，链与链间的相邻点也要考虑清楚。

#include <algorithm>
#include <iterator>
#include <iostream>
#include <cstring>
#include <cstdlib>
#include <iomanip>
#include <bitset>
#include <cctype>
#include <cstdio>
#include <string>
#include <vector>
#include <stack>
#include <cmath>
#include <queue>
#include <list>
#include <map>
#include <set>
#include <cassert>

using namespace std;
#define lson (l, mid, rt << 1)
#define rson (mid + 1, r, rt << 1 | 1)
#define debug(x) cerr << #x << " = " << x << "\n";
#define pb push_back
#define pq priority_queue

typedef long long ll;
typedef unsigned long long ull;
//typedef __int128 bll;
typedef pair<ll, ll> pll;
typedef pair<int, int> pii;
typedef pair<int, pii> p3;

//priority_queue<int> q;//这是一个大根堆q
//priority_queue<int,vector<int>,greater<int> >q;//这是一个小根堆q
#define fi first
#define se second
//#define endl '\n'

#define boost                    \
    ios::sync_with_stdio(false); \
    cin.tie(0)
#define rep(a, b, c) for (int a = (b); a <= (c); ++a)
#define max3(a, b, c) max(max(a, b), c);
#define min3(a, b, c) min(min(a, b), c);

const ll oo = 1ll << 17;
const ll mos = 0x7FFFFFFF;  //2147483647
const ll nmos = 0x80000000; //-2147483648
const int inf = 0x3f3f3f3f;
const ll inff = 0x3f3f3f3f3f3f3f3f; //18
const ll mod = 2147483648;
const double esp = 1e-8;
const double PI = acos(-1.0);
const double PHI = 0.61803399; //黄金分割点
const double tPHI = 0.38196601;

template <typename T>
inline T read(T &x) {
    x = 0;
    int f = 0;
    char ch = getchar();
    while (ch < '0' || ch > '9') f |= (ch == '-'), ch = getchar();
    while (ch >= '0' && ch <= '9') x = x * 10 + ch - '0', ch = getchar();
    return x = f ? -x : x;
}

inline void cmax(int &x, int y) {
    if (x < y) x = y;
}
inline void cmax(ll &x, ll y) {
    if (x < y) x = y;
}
inline void cmin(int &x, int y) {
    if (x > y) x = y;
}
inline void cmin(ll &x, ll y) {
    if (x > y) x = y;
}

/*-----------------------showtime----------------------*/

const int maxn = 1e5 + 9;
int a[maxn], b[maxn];
int dp[maxn], sz[maxn], fa[maxn], son[maxn];
vector<int> mp[maxn];

void dfs1(int u, int f, int deep) {
    dp[u] = deep;
    fa[u] = f;
    sz[u] = 1;
    int mx = 0;
    for (int i = 0; i < mp[u].size(); i++) {
        int v = mp[u][i];
        if (v == f) continue;
        dfs1(v, u, deep + 1);
        sz[u] += sz[v];
        if (sz[v] > mx) { mx = sz[v], son[u] = v; }
    }
}

int top[maxn], id[maxn], cnt = 0;
void dfs2(int u, int f, int topf) {
    top[u] = topf;
    id[u] = ++cnt;
    b[cnt] = a[u];
    if (son[u]) dfs2(son[u], u, topf);
    for (int i = 0; i < mp[u].size(); i++) {
        int v = mp[u][i];
        if (v == f || v == son[u]) continue;
        dfs2(v, u, v);
    }
}

int tag[maxn << 2], lazy[maxn << 2], ly[maxn << 2], rz[maxn << 2];
void pushup(int rt) {
    tag[rt] = tag[rt << 1] + tag[rt << 1 | 1];
    rz[rt] = rz[rt << 1];
    ly[rt] = ly[rt << 1 | 1];
    if (ly[rt << 1] == rz[rt << 1 | 1]) tag[rt]--;
}

void build(int l, int r, int rt) {
    if (l == r) {
        tag[rt] = 1;
        ly[rt] = rz[rt] = b[l];
        return;
    }
    int mid = (l + r) >> 1;
    build(l, mid, rt << 1);
    build(mid + 1, r, rt << 1 | 1);
    pushup(rt);
    // cout<<l<<" "<<r<<" "<<rz[rt] << "  " << ly[rt]<<endl;
}

void pushdown(int l, int r, int rt) {
    tag[rt << 1] = tag[rt << 1 | 1] = 1;
    ly[rt << 1] = rz[rt << 1] = lazy[rt];
    ly[rt << 1 | 1] = rz[rt << 1 | 1] = lazy[rt];
    lazy[rt << 1] = lazy[rt << 1 | 1] = lazy[rt];
    lazy[rt] = 0;
}

void update(int L, int R, int c, int l, int r, int rt) {
    if (l >= L && r <= R) {
        tag[rt] = 1;
        ly[rt] = rz[rt] = c;
        lazy[rt] = c;
        return;
    }
    int mid = (l + r) >> 1;
    if (lazy[rt]) pushdown(l, r, rt);
    if (mid >= L) update(L, R, c, l, mid, rt << 1);
    if (mid < R) update(L, R, c, mid + 1, r, rt << 1 | 1);
    pushup(rt);
}

int n, m;

void solve(int x, int y, int c) {
    while (top[x] != top[y]) {
        if (dp[top[x]] < dp[top[y]]) swap(x, y);
        update(id[top[x]], id[x], c, 1, n, 1);
        x = fa[top[x]];
    }
    if (dp[x] > dp[y]) swap(x, y);
    update(id[x], id[y], c, 1, n, 1);
}

int query(int L, int R, int l, int r, int rt, int &tmpl, int &tmpr) {
    if (l >= L && r <= R) {
        if (l == L) tmpl = rz[rt];
        if (r == R) tmpr = ly[rt];

        return tag[rt];
    }
    int mid = (l + r) >> 1;
    if (lazy[rt]) pushdown(l, r, rt);
    int res = 0;
    int ql = -1, qr = -1;
    if (mid >= L) {
        res += query(L, R, l, mid, rt << 1, tmpl, tmpr);
        ql = ly[rt << 1];
    }
    if (mid < R) {
        res += query(L, R, mid + 1, r, rt << 1 | 1, tmpl, tmpr);
        qr = rz[rt << 1 | 1];
    }
    if (ql == qr && ql != -1) res--;
    pushup(rt);
    return res;
}

int cal(int x, int y) {
    int res = 0, lax = -1, lay = -1;

    while (top[x] != top[y]) {
        int tmpl = -1, tmpr = -1;
        if (dp[top[x]] > dp[top[y]]) {
            res += query(id[top[x]], id[x], 1, n, 1, tmpl, tmpr); //tmp,flag 1,r
            if (lax == tmpr) res--;
            lax = tmpl;
            x = fa[top[x]];
        } else {
            res += query(id[top[y]], id[y], 1, n, 1, tmpl, tmpr);
            if (lay == tmpr) res--;
            lay = tmpl;
            y = fa[top[y]];
        }

        // cout<<tmpl<<" !! "<<tmpr<<endl;
    }

    if (dp[x] > dp[y]) {
        int tmpl = -1, tmpr = -1;
        res += query(id[y], id[x], 1, n, 1, tmpl, tmpr);
        if (tmpr == lax) res--;
        if (tmpl == lay) res--;
    } else {
        int tmpl = -1, tmpr = -1;
        res += query(id[x], id[y], 1, n, 1, tmpl, tmpr);
        if (tmpl == lax) res--;
        if (tmpr == lay) res--;
    }

    return res;
}

int main() {
    scanf("%d%d", &n, &m);
    rep(i, 1, n) scanf("%d", &a[i]);
    rep(i, 1, n - 1) {
        int u, v;
        scanf("%d%d", &u, &v);
        mp[u].pb(v);
        mp[v].pb(u);
    }
    dfs1(1, 1, 1);
    dfs2(1, 1, 1);
    build(1, n, 1);

    while (m--) {
        char str[5];
        scanf("%s", str);
        if (str[0] == 'C') {
            int x, y, z;
            scanf("%d%d%d", &x, &y, &z);
            solve(x, y, z);
        } else {
            int x, y;
            scanf("%d%d", &x, &y);
            printf("%d\n", cal(x, y));
        }
    }
    return 0;
}