BZOJ 2243

2243: [SDOI2011]染色

树链剖分+线段树
题意是多少段，不是多少种。
所以维护好left color和right color就好了

BZOJ 2243

/**************************************************************
    Problem: 2243
    User: Dream_Tonight
    Language: C++
    Result: Accepted
    Time:6012 ms
    Memory:17404 kb
****************************************************************/
 
#include <algorithm>
#include <iostream>
#include <cstring>
#include <vector>
#include <string>
#include <cstdio>
#include <cmath>
#include <stack>
#include <queue>
#include <list>
#include <map>
#include <set>
 
using namespace std;
#define ls rt << 1
#define rs rt << 1 | 1
#define lson l , mid , ls
#define rson mid + 1 , r , rs
typedef long long ll;
typedef pair<int, int> pii;
const int maxn = 1e5 + 7;
const int inf = 2139062143;
const int mod = 10007;
 
int fa[maxn], son[maxn], sz[maxn], rk[maxn];
int top[maxn], head[maxn], deep[maxn], id[maxn];
int color[maxn], sum[maxn << 2], lc[maxn << 2], rc[maxn << 2], tag[maxn << 2];
int cnt = 0, ret = 0;
int a, b, c, LC, RC;
char ch;
 
struct node {
    int to, next;
} edge[maxn << 1];
 
void add(int u, int v) {
    edge[++ret] = (node) {v, head[u]}, head[u] = ret;
    edge[++ret] = (node) {u, head[v]}, head[v] = ret;
}
 
void dfs1(int x, int f) {
    sz[x] = 1, fa[x] = f, deep[x] = deep[f] + 1, son[x] = 0;
    for (int i = head[x]; i; i = edge[i].next) {
        int v = edge[i].to;
        if (v == f) continue;
        dfs1(v, x);
        if (sz[son[x]] < sz[v]) son[x] = v;
        sz[x] += sz[v];
    }
}
 
void dfs2(int x, int tp) {
    top[x] = tp, id[x] = ++cnt, rk[cnt] = color[x];
    if (son[x]) dfs2(son[x], tp);
    for (int i = head[x]; i; i = edge[i].next) {
        int v = edge[i].to;
        if (v == son[x] || v == fa[x]) continue;
        dfs2(v, v);
    }
}
 
void push_up(int rt) {
    sum[rt] = sum[ls] + sum[rs];
    if (rc[ls] == lc[rs]) sum[rt]--;
    lc[rt] = lc[ls], rc[rt] = rc[rs];
}
 
void push_down(int rt) {
    if (tag[rt] == -1) return;
    sum[ls] = sum[rs] = 1;
    lc[ls] = rc[ls] = tag[rt];
    lc[rs] = rc[rs] = tag[rt];
    tag[ls] = tag[rs] = tag[rt];
    tag[rt] = -1;
}
 
void build(int l, int r, int rt) {
    if (l == r) {
        sum[rt] = 1;
        lc[rt] = rc[rt] = rk[l];
        return;
    }
    int mid = (l + r) / 2;
    build(lson);
    build(rson);
    push_up(rt);
}
 
int query(int l, int r, int rt, int L, int R) {
    if (L <= l && r <= R) {
        if (l == L) LC = lc[rt];
        if (r == R) RC = rc[rt];
        return sum[rt];
    }
    push_down(rt);
    int mid = (l + r) / 2, res = 0;
    if (mid >= L) res += query(lson, L, R);
    if (mid < R) res += query(rson, L, R);
    if (L <= mid && mid < R) res -= rc[ls] == lc[rs];
    push_up(rt);
    return res;
}
 
void update(int l, int r, int rt, int L, int R, int V) {
    if (L <= l && r <= R) {
        tag[rt] = lc[rt] = rc[rt] = V;
        sum[rt] = 1;
        return;
    }
    push_down(rt);
    int mid = (l + r) / 2;
    if (mid >= L) update(lson, L, R, V);
    if (mid < R) update(rson, L, R, V);
    push_up(rt);
}
 
void update_range(int l, int r, int col, int n) {
    while (top[l] != top[r]) {
        if (deep[top[l]] < deep[top[r]]) swap(l, r);
        update(1, n, 1, id[top[l]], id[l], col);
        l = fa[top[l]];
    }
    if (deep[l] > deep[r]) swap(l, r);
    update(1, n, 1, id[l], id[r], col);
}
 
int query_range(int l, int r, int n) {
    int res = 0, ll = -1, rr = -1;
    while (top[l] != top[r]) {
        if (deep[top[l]] < deep[top[r]]) swap(l, r), swap(ll, rr);
        res += query(1, n, 1, id[top[l]], id[l]);
        if (RC == ll) res--;
        ll = LC;
        l = fa[top[l]];
    }
    if (deep[l] > deep[r]) swap(l, r), swap(ll, rr);
    res += query(1, n, 1, id[l], id[r]);
    if (LC == ll) res--;
    if (RC == rr) res--;
    return res;
}
 
int main() {
#ifndef ONLINE_JUDGE
    freopen("weather.in", "r", stdin);
#endif
    int x, y, n, m;
    scanf("%d%d", &n, &m);
    memset(tag, -1, sizeof(tag));
    for (int i = 1; i <= n; i++)scanf("%d", &color[i]);
    for (int i = 1; i < n; i++) {
        scanf("%d%d", &x, &y);
        add(x, y);
    }
    dfs1(1, 0), dfs2(1, 0);
    build(1, n, 1);
    for (int i = 1; i <= m; i++) {
        scanf(" %c", &ch);
        if (ch == 'Q') {
            scanf("%d%d", &a, &b);
            printf("%d\n", query_range(a, b, n));
        } else {
            scanf("%d%d%d", &a, &b, &c);
            update_range(a, b, c, n);
        }
    }
    return 0;
}